242 lines
8.8 KiB
Plaintext
242 lines
8.8 KiB
Plaintext
/*
|
||
* fused_rmsnorm_amax_quantize.cu
|
||
*
|
||
* Fused RMSNorm + amax + NVFP4 quantize.
|
||
* Combines 5+ kernel launches into 2:
|
||
* Kernel 1: compute RMS + row-wise amax → gsa (GPU buffer)
|
||
* Kernel 2: normalize + quantize using gsa from GPU buffer
|
||
*
|
||
* Eliminates:
|
||
* - 3-4 BF16 RMSNorm kernel launches (x^2, mean, rsqrt, multiply)
|
||
* - 1 amax kernel launch
|
||
* - Total: ~4-5 launches per RMSNorm site × ~366 sites = ~1500 launches/token
|
||
*
|
||
* Production: used for attention input norm, FFN input norm, q_a_norm, kv_norm.
|
||
*/
|
||
|
||
#include <cuda.h>
|
||
#include <cuda_bf16.h>
|
||
#include <cuda_fp8.h>
|
||
#include <cfloat>
|
||
#include <cmath>
|
||
|
||
// FP4 E2M1 look-up table (same as production)
|
||
__device__ __constant__ float FP4_LUT[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
|
||
|
||
// ============================================================================
|
||
// Kernel 1: Compute RMS + row-wise amax + gsa
|
||
// ============================================================================
|
||
// Each block handles one row. Threadblock: (WARP_SIZE, 1, 1) per row.
|
||
// For hidden_size=7168: 7168/32 = 224 warps needed. Use multiple blocks.
|
||
|
||
__global__ void rmsnorm_amax_gsa_kernel(
|
||
const __nv_bfloat16* __restrict__ x, // (M, N) BF16 row-major
|
||
const float* __restrict__ norm_weight, // (N,) FP32
|
||
float* __restrict__ gsa_out, // (M,) FP32 — per-row gsa
|
||
float* __restrict__ rms_out, // (M,) FP32 — per-row RMS (for kernel 2)
|
||
const int M,
|
||
const int N,
|
||
const float eps,
|
||
const float divisor // gsa = amax / divisor
|
||
) {
|
||
const int row = blockIdx.x;
|
||
if (row >= M) return;
|
||
|
||
const __nv_bfloat16* x_row = x + row * N;
|
||
|
||
// Step 1: Compute RMS (sum of x^2) using warp reduction
|
||
float sum_sq = 0.0f;
|
||
float row_amax = 0.0f;
|
||
|
||
for (int col = threadIdx.x; col < N; col += blockDim.x) {
|
||
float val = __bfloat162float(x_row[col]);
|
||
float val_sq = val * val;
|
||
sum_sq += val_sq;
|
||
float abs_val = fabsf(val);
|
||
if (abs_val > row_amax) row_amax = abs_val;
|
||
}
|
||
|
||
// Warp-level reduction for sum_sq and amax
|
||
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
|
||
sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset);
|
||
row_amax = fmaxf(row_amax, __shfl_down_sync(0xFFFFFFFF, row_amax, offset));
|
||
}
|
||
|
||
// Block-level reduction (if blockDim.x > warpSize)
|
||
// Use shared memory for cross-warp reduction
|
||
__shared__ float s_sum_sq[32]; // max 32 warps
|
||
__shared__ float s_amax[32];
|
||
|
||
int lane = threadIdx.x % warpSize;
|
||
int warp_id = threadIdx.x / warpSize;
|
||
int num_warps = blockDim.x / warpSize;
|
||
|
||
if (lane == 0) {
|
||
s_sum_sq[warp_id] = sum_sq;
|
||
s_amax[warp_id] = row_amax;
|
||
}
|
||
__syncthreads();
|
||
|
||
if (warp_id == 0) {
|
||
sum_sq = (lane < num_warps) ? s_sum_sq[lane] : 0.0f;
|
||
row_amax = (lane < num_warps) ? s_amax[lane] : 0.0f;
|
||
for (int offset = warpSize / 2; offset > 0; offset /= 2) {
|
||
sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset);
|
||
row_amax = fmaxf(row_amax, __shfl_down_sync(0xFFFFFFFF, row_amax, offset));
|
||
}
|
||
}
|
||
|
||
if (threadIdx.x == 0) {
|
||
float rms = sqrtf(sum_sq / N + eps);
|
||
float gsa = row_amax / divisor;
|
||
gsa_out[row] = gsa;
|
||
rms_out[row] = rms;
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Kernel 2: RMSNorm + quantize using gsa from GPU buffer
|
||
// ============================================================================
|
||
// Each block handles one row. Quantizes x * rsqrt(rms) * weight to NVFP4.
|
||
|
||
__global__ void rmsnorm_quantize_nvfp4_kernel(
|
||
const __nv_bfloat16* __restrict__ x, // (M, N) BF16 row-major
|
||
const float* __restrict__ norm_weight, // (N,) FP32
|
||
const float* __restrict__ gsa, // (M,) FP32 — per-row global scale
|
||
const float* __restrict__ rms, // (M,) FP32 — per-row RMS
|
||
uint8_t* __restrict__ x_fp4, // (M, N//2) FP4 packed
|
||
__nv_fp8_e4m3* __restrict__ x_sf, // (M, N//16) E4M3 block scales
|
||
const int M,
|
||
const int N,
|
||
const float eps
|
||
) {
|
||
const int row = blockIdx.x;
|
||
if (row >= M) return;
|
||
|
||
const __nv_bfloat16* x_row = x + row * N;
|
||
float row_gsa = gsa[row];
|
||
float row_rms = rms[row];
|
||
float inv_rms = 1.0f / row_rms; // rsqrt(sum/N + eps) = 1/rms
|
||
|
||
// Process 16 elements per thread (one microblock)
|
||
// Each thread writes 8 FP4 values (16 BF16 → 8 bytes packed)
|
||
// Block scales: 1 E4M3 per 16 BF16 values
|
||
|
||
const int num_blocks = N / 16; // number of 16-element microblocks
|
||
const int blocks_per_thread = 1;
|
||
const int total_threads = blockDim.x;
|
||
|
||
for (int blk = threadIdx.x; blk < num_blocks; blk += total_threads) {
|
||
int col = blk * 16;
|
||
|
||
// Load 16 BF16 values, normalize, and compute block scale
|
||
float vals[16];
|
||
float block_amax = 0.0f;
|
||
|
||
for (int i = 0; i < 16; i++) {
|
||
float v = __bfloat162float(x_row[col + i]);
|
||
v = v * inv_rms * norm_weight[col + i]; // RMSNorm
|
||
vals[i] = v;
|
||
float av = fabsf(v);
|
||
if (av > block_amax) block_amax = av;
|
||
}
|
||
|
||
// Compute block scale
|
||
// block_scale = amax / (gsa * 6.0)
|
||
// E4M3 range is [-448, 448], FP4 range is [-6, 6]
|
||
// Quantized value = val / (gsa * block_scale)
|
||
float block_scale = block_amax / (row_gsa * 6.0f);
|
||
|
||
// Clamp block_scale to E4M3 range
|
||
if (block_scale > 448.0f) block_scale = 448.0f;
|
||
if (block_scale < 0.0f) block_scale = 0.0f;
|
||
|
||
// Convert block_scale to E4M3
|
||
__nv_fp8_e4m3 bs_e4m3;
|
||
bs_e4m3.__x = 0;
|
||
if (block_scale > 0.0f) {
|
||
// Simple float → E4M3 conversion
|
||
// E4M3: 1 sign + 4 exponent + 3 mantissa bits
|
||
// Range: 2^-9 * (1 + 0/8) to 2^4 * (1 + 7/8) ≈ 0.00195 to 30.0
|
||
// But we use it as a scale factor, so range [0, 448]
|
||
// Actually, E4M3 max = 448.0, and we clamp to that above.
|
||
// Use CUDA's native conversion:
|
||
float bs_f = block_scale;
|
||
__nv_fp8_e4m3 tmp;
|
||
tmp = __nv_fp8_e4m3(bs_f);
|
||
bs_e4m3 = tmp;
|
||
}
|
||
|
||
// Quantize 16 values to FP4 (8 bytes)
|
||
uint8_t packed = 0;
|
||
|
||
for (int i = 0; i < 16; i += 2) {
|
||
// Two values per byte
|
||
float v0 = vals[i] / (row_gsa * block_scale);
|
||
float v1 = vals[i + 1] / (row_gsa * block_scale);
|
||
|
||
// Clamp to FP4 range [-6, 6]
|
||
v0 = fmaxf(fminf(v0, 6.0f), -6.0f);
|
||
v1 = fmaxf(fminf(v1, 6.0f), -6.0f);
|
||
|
||
// Find nearest FP4 value
|
||
auto quantize_fp4 = [](float v) -> uint8_t {
|
||
bool sign = v < 0;
|
||
float av = fabsf(v);
|
||
float best_err = 1e10f;
|
||
uint8_t best_idx = 0;
|
||
for (int k = 0; k < 8; k++) {
|
||
float err = fabsf(av - FP4_LUT[k]);
|
||
if (err < best_err) {
|
||
best_err = err;
|
||
best_idx = k;
|
||
}
|
||
}
|
||
return best_idx | (sign ? 0x08 : 0x00);
|
||
};
|
||
|
||
uint8_t lo = quantize_fp4(v0);
|
||
uint8_t hi = quantize_fp4(v1);
|
||
packed |= (lo | (hi << 4)) << (i / 2 * 8); // pack 2 nibbles per byte
|
||
}
|
||
|
||
// Wait, this packing is wrong. Each byte packs 2 FP4 values (nibbles).
|
||
// 16 BF16 → 8 FP4 bytes → 16 FP4 nibbles.
|
||
// Actually: 16 BF16 values → 8 bytes, each byte has 2 FP4 nibbles.
|
||
|
||
// Let me redo this properly:
|
||
uint8_t fp4_bytes[8] = {0};
|
||
for (int i = 0; i < 16; i += 2) {
|
||
float v0 = vals[i] / (row_gsa * block_scale);
|
||
float v1 = vals[i + 1] / (row_gsa * block_scale);
|
||
v0 = fmaxf(fminf(v0, 6.0f), -6.0f);
|
||
v1 = fmaxf(fminf(v1, 6.0f), -6.0f);
|
||
|
||
auto qf4 = [](float v) -> uint8_t {
|
||
bool sign = v < 0;
|
||
float av = fabsf(v);
|
||
float best_err = 1e10f;
|
||
uint8_t best_idx = 0;
|
||
for (int k = 0; k < 8; k++) {
|
||
float err = fabsf(av - FP4_LUT[k]);
|
||
if (err < best_err) { best_err = err; best_idx = k; }
|
||
}
|
||
return best_idx | (sign ? 0x08 : 0x00);
|
||
};
|
||
|
||
uint8_t lo = qf4(v0);
|
||
uint8_t hi = qf4(v1);
|
||
fp4_bytes[i / 2] = lo | (hi << 4);
|
||
}
|
||
|
||
// Write 8 FP4 bytes
|
||
uint8_t* dst = x_fp4 + row * (N / 2) + blk * 8;
|
||
for (int i = 0; i < 8; i++) {
|
||
dst[i] = fp4_bytes[i];
|
||
}
|
||
|
||
// Write block scale
|
||
x_sf[row * (N / 16) + blk] = bs_e4m3;
|
||
}
|
||
}
|