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

P4: Fix fused RMSNorm kernel — match quantize_nvfp4.cu encoding

Browse Source 
- Use half_step_to_e2m1 for E2M1 FP4 quantization (not LUT search)
- Use __nv_fp8_e4m3 + memcpy for block scale (not reinterpret_cast)
- Pack nibbles as (nibbles[2*i+1] << 4) | nibbles[2*i] (same as prod)
- Output uint8 buffers, then .view() to FP4/FP8 dtypes
- Handle near-zero block scale same as quantize_nvfp4.cu
This commit is contained in:
biondizzle
2026-06-02 16:28:44 +00:00
parent 794ebaf7e5
commit 29f836d711
1 changed files with 58 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
dsv4/kernels/cuda/fused_rmsnorm_quantize.cu
View File

@@ -14,41 +14,42 @@
* Kernel 2: rmsnorm_quantize_nvfp4_kernel
* - Read gsa + inv_rms from GPU buffers (no CPU sync)
* - Normalize: val = x * inv_rms * weight
* - Quantize to NVFP4 using the gsa from Kernel 1
* - Quantize to NVFP4 using the same proven path as quantize_nvfp4.cu
* - Write FP4 data + E4M3 block scales
*
* Usage sites per decode step:
* - 2 RMSNorm per layer (attn_norm + ffn_norm) × 61 layers = 122 calls
* - Each currently: 4+ launches (rmsnorm) + 2 launches (amax+quantize) = 6+
* - After fusion: 2 launches per site → 244 launches eliminated
*
* Grid strategy:
* Kernel 1: 1 block per row (M blocks, 1 threadblock per row)
* - Each block computes rms + amax for its row using warp reductions
* Kernel 2: 1 block per (row, 16-element microblock)
* - Same grid as quantize_nvfp4_kernel for consistency
* Quantization is bit-identical to quantize_nvfp4.cu:
* - half_step_to_e2m1 for E2M1 encoding
* - __nv_fp8_e4m3 for block scale
* - (nibbles[2*i+1] << 4) | nibbles[2*i] packing
*/
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cuda_fp8.hpp>
#include <ATen/ATen.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include <cstdint>
#include <cfloat>
#include <cmath>
#include <cstring>
// FP4 E2M1 look-up table — same as production quantize_nvfp4.cu
__device__ __constant__ float FP4_LUT[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
// FP4 E2M1 half-step → index mapping (same as quantize_nvfp4.cu)
__device__ __forceinline__ int half_step_to_e2m1(int hs) {
if (hs <= 4) return hs;
if (hs <= 5) return 4;
if (hs <= 7) return 5;
if (hs <= 10) return 6;
return 7;
}
// ============================================================================
// Kernel 1: Compute RMS + amax of normalized output → gsa per row
// ============================================================================
// Each block processes one row of (M, N).
// Threadblock: blockDim.x threads per row (must be multiple of warpSize).
// For hidden_size=7168: 7168/32 = 224 threads. Use blockDim.x=256 for alignment.
__global__ void rmsnorm_amax_gsa_kernel(
const __nv_bfloat16* __restrict__ x, // (M, N) BF16 row-major
@@ -65,13 +66,7 @@ __global__ void rmsnorm_amax_gsa_kernel(
const __nv_bfloat16* x_row = x + (size_t)row * N;
// Step 1: Compute sum(x^2) and amax of (x * inv_rms * weight) in one pass
// We need both sum_sq (for RMS) and the amax of the NORMALIZED output
// Can't compute normalized amax without RMS yet, so we do it in two sub-passes:
// Sub-pass 1: compute sum(x^2) for RMS
// Sub-pass 2: after RMS known, compute amax of normalized output
// Sub-pass 1: sum of squares
// Sub-pass 1: compute sum(x^2) for RMS
float sum_sq = 0.0f;
for (int col = threadIdx.x; col < N; col += blockDim.x) {
float val = __bfloat162float(x_row[col]);
@@ -107,7 +102,7 @@ __global__ void rmsnorm_amax_gsa_kernel(
__shared__ float s_inv_rms;
if (threadIdx.x == 0) {
float rms = sqrtf(row_sum_sq / N + eps);
s_inv_rms = 1.0f / rms;
s_inv_rms = 1.0f / fmaxf(rms, 1e-8f);
}
__syncthreads();
float inv_rms = s_inv_rms;
@@ -149,23 +144,15 @@ __global__ void rmsnorm_amax_gsa_kernel(
// ============================================================================
// Same grid as quantize_nvfp4_kernel: (N/16, M, 1)
// Each CTA processes one 16-element microblock in one row.
__device__ __forceinline__ int half_step_to_e2m1(int hs) {
// Maps half-step indices to E2M1 indices (same as quantize_nvfp4.cu)
if (hs <= 4) return hs;
if (hs <= 5) return 4;
if (hs <= 7) return 5;
if (hs <= 10) return 6;
return 7;
}
// Bit-identical quantization to quantize_nvfp4.cu.
__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__ inv_rms, // (M,) FP32 — per-row 1/rms
uint8_t* __restrict__ x_fp4, // (M, N//2) FP4 packed
uint8_t* __restrict__ x_sf, // (M, N//16) E4M3 block scales
uint8_t* __restrict__ out_fp4, // (M, N//2) FP4 packed
uint8_t* __restrict__ out_sf, // (M, N//16) E4M3 block scales (uint8 view)
const int M,
const int N
) {
@@ -178,7 +165,7 @@ __global__ void rmsnorm_quantize_nvfp4_kernel(
float row_gsa = gsa[row];
float row_inv_rms = inv_rms[row];
// Step 1: Load 16 BF16 elements, normalize, compute block amax
// Step 1: Load 16 BF16 elements, normalize (RMSNorm), compute block amax
float vals[16];
float block_amax = 0.0f;
const int col_base = n_block * 16;
@@ -196,56 +183,38 @@ __global__ void rmsnorm_quantize_nvfp4_kernel(
}
}
// Step 2: Compute block scale
float block_scale = block_amax / (row_gsa * 6.0f);
// Clamp to E4M3 representable range
block_scale = fmaxf(block_scale, 0.0f);
block_scale = fminf(block_scale, 448.0f);
// Step 2: Compute FP8 E4M3 block scale (same as quantize_nvfp4.cu)
// block_scale = block_amax / (gsa * 6.0)
float bsf = block_amax / (row_gsa * 6.0f);
if (block_amax < row_gsa * 6.0f * 0.001953125f) {
bsf = 0.0f;
for (int i = 0; i < 16; i++) vals[i] = 0.0f;
}
__nv_fp8_e4m3 bsf8_obj(bsf);
float bs = (float)bsf8_obj; // dequantized block scale for FP4 computation
uint8_t bsf8;
memcpy(&bsf8, &bsf8_obj, 1);
// Convert block_scale to E4M3
__nv_fp8_e4m3 bs_e4m3;
bs_e4m3 = __nv_fp8_e4m3(block_scale);
// Step 3: Quantize 16 values to FP4
// Same proven quantization path as quantize_nvfp4.cu
uint8_t fp4_bytes[8];
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);
// Clamp to FP4 range [-6, 6]
v0 = fmaxf(fminf(v0, 6.0f), -6.0f);
v1 = fmaxf(fminf(v1, 6.0f), -6.0f);
// Quantize using LUT — find nearest FP4 value
// lo nibble: v0, hi nibble: v1
uint8_t lo_idx = 0, hi_idx = 0;
float lo_err = 1e10f, hi_err = 1e10f;
for (int k = 0; k < 8; k++) {
float e0 = fabsf(fabsf(v0) - FP4_LUT[k]);
float e1 = fabsf(fabsf(v1) - FP4_LUT[k]);
if (e0 < lo_err) { lo_err = e0; lo_idx = k; }
if (e1 < hi_err) { hi_err = e1; hi_idx = k; }
}
// Apply sign: bit 3 = sign
if (v0 < 0) lo_idx |= 0x08;
if (v1 < 0) hi_idx |= 0x08;
fp4_bytes[i / 2] = lo_idx | (hi_idx << 4);
// Step 3: Quantize each value to FP4 E2M1 (same as quantize_nvfp4.cu)
uint8_t nibbles[16];
for (int i = 0; i < 16; i++) {
if (bs < 1e-8f) { nibbles[i] = 0; continue; }
float s = vals[i] / (row_gsa * bs); // scale by gsa * block_scale
int hs = __float2int_rn(fminf(fabsf(s), 6.0f) * 2.0f);
if (hs > 12) hs = 12;
int idx = half_step_to_e2m1(hs);
if (s < 0) idx += 8;
nibbles[i] = idx;
}
// Step 4: Write FP4 data and block scale
// FP4 data: (M, N//2) — row-major, 8 bytes per 16-element microblock
uint8_t* dst_fp4 = x_fp4 + (size_t)row * (N / 2) + n_block * 8;
// Step 4: Pack pairs: (nibbles[2*i+1] << 4) | nibbles[2*i] (same as quantize_nvfp4.cu)
for (int i = 0; i < 8; i++) {
dst_fp4[i] = fp4_bytes[i];
out_fp4[(size_t)row * (N / 2) + n_block * 8 + i] =
(nibbles[2 * i + 1] << 4) | nibbles[2 * i];
}
// Block scale: (M, N//16) — one E4M3 per microblock
// x_sf is uint8_t view of E4M3 data
__nv_fp8_e4m3* dst_sf = reinterpret_cast<__nv_fp8_e4m3*>(
x_sf + (size_t)row * (N / 16) * sizeof(__nv_fp8_e4m3) + n_block * sizeof(__nv_fp8_e4m3));
*dst_sf = bs_e4m3;
// Step 5: Write FP8 block scale (uint8 view, same as quantize_nvfp4.cu)
out_sf[(size_t)row * (N / 16) + n_block] = bsf8;
}
// ============================================================================
@@ -270,11 +239,11 @@ rmsnorm_quantize_nvfp4_cuda(
auto stream = c10::cuda::getCurrentCUDAStream();
auto options = x.options();
// Output buffers
// Output buffers (uint8, then .view() to FP4/FP8 dtypes)
auto gsa = torch::empty({M}, options.dtype(torch::kFloat32));
auto inv_rms = torch::empty({M}, options.dtype(torch::kFloat32));
auto x_fp4 = torch::empty({M, N / 2}, options.dtype(torch::kUInt8));
auto x_sf = torch::empty({M, N / 16}, options.dtype(torch::kFloat8E4M3FN));
auto x_sf = torch::empty({M, N / 16}, options.dtype(torch::kUInt8));
// Kernel 1: RMSNorm + amax → gsa (1 block per row)
const int threads1 = 256; // 8 warps, handles up to N=8192
@@ -296,11 +265,17 @@ rmsnorm_quantize_nvfp4_cuda(
gsa.data_ptr<float>(),
inv_rms.data_ptr<float>(),
x_fp4.data_ptr<uint8_t>(),
reinterpret_cast<uint8_t*>(x_sf.data_ptr()),
x_sf.data_ptr<uint8_t>(),
M, N
);
return std::make_tuple(x_fp4, x_sf, gsa, inv_rms);
// View as proper dtypes (same as quantize_nvfp4.cu)
return std::make_tuple(
x_fp4.view(torch::kFloat4_e2m1fn_x2),
x_sf.view(torch::kFloat8_e4m3fn),
gsa,
inv_rms
);
}
// Standalone kernel 1 entry point (for testing / when only gsa needed)
@@ -311,8 +286,7 @@ torch::Tensor rmsnorm_amax_gsa_cuda(
double divisor
) {
TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
TORCH_CHECK(x.scalar_type() == torch::kFloat16 || x.scalar_type() == torch::kBFloat16,
"x must be BF16");
TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be BF16");
const int M = x.size(0);
const int N = x.size(1);