diff --git a/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu index 2af431b5..e689eea0 100644 --- a/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu +++ b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu @@ -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 #include #include #include +#include #include #include #include #include #include #include +#include -// 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(), inv_rms.data_ptr(), x_fp4.data_ptr(), - reinterpret_cast(x_sf.data_ptr()), + x_sf.data_ptr(), 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);