biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P3 WIP: fused RMSNorm + quantize kernel skeleton (not yet integrated)

Browse Source
This commit is contained in:
biondizzle
2026-06-02 09:02:52 +00:00
parent b13c1057f5
commit 851ec9b4d5
1 changed files with 241 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

241
dsv4/kernels/cuda/fused_rmsnorm_quantize.cu Normal file
View File

@@ -0,0 +1,241 @@
/*
* 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;
}
}