diff --git a/PERFORMANCE_AUDIT.md b/PERFORMANCE_AUDIT.md index 7ecb1d7e..db964143 100644 --- a/PERFORMANCE_AUDIT.md +++ b/PERFORMANCE_AUDIT.md @@ -212,97 +212,6 @@ Blockers in v17: **Sequence:** land after P0/P1/P2/P3 so the captured graph reflects the post-fusion structure. ---- - -# PART 4 — TURBOQUANT: ARCHITECTURAL VERDICT - -Reading `turboquant/`: this is an **ICLR 2026 paper implementation** of -vector-quantization KV compression. Two algorithms: -- MSE-quantize keys/values via codebook (3 bit by default) -- Inner-product-aware quantize keys (preserves dot products) via Algorithm 2 -- Per-vector L2-norm preserved separately, plus QJL sign sketch for - residual recovery - -Operational shape: -- Operates on **standard MHA/GQA shape** `(..., n_heads, head_dim)`, - head_dim typically 128. -- Requires a `head_dim × head_dim` rotation matrix per layer (precomputed - from random seed, shared across heads). -- Has a Triton fused-decode kernel that computes attention scores directly - from packed codebook indices. -- vLLM integration via `turboquant/vllm_attn_backend.py`. - -## Why it doesn't fit DSv4 - -Three structural mismatches, in order of severity: - -### 1. The DSv4 KV cache is already a learned compression - -DSv4 doesn't store per-token KV. The CSA compressor's whole job is to -reduce m=4 tokens into 1 compressed entry via a softmax-weighted mix. -That entry is what gets cached. TurboQuant quantizes the *post-projection -per-token KV* of standard attention — exactly the thing DSv4 has -already replaced with a learned compressor. **You'd be applying a lossy -compression on top of an already-lossy compression**, which (a) compounds -loss in an uncontrolled way and (b) attacks the wrong dimension. The -compressed entries are already 4× (CSA) or 128× (HCA) reduced in the -sequence dimension; further reducing the *head dimension* via codebook -gives little additional savings (you're already attending over very few -entries per query) at high quality cost. - -### 2. Wrong shape, wrong primitive - -TurboQuant operates on `(..., n_heads, head_dim=128)` per-token vectors -and uses a `128×128` random rotation. DSv4's compressed cache is shape -`(n_comp, head_dim=512)` — no head dimension. The whole "rotate the head -dim" abstraction needs to be reworked, and once you do, you're writing -new code that isn't TurboQuant anymore. - -For the indexer keys, the storage *is* per-block 128-dim, which is closer -to TurboQuant's natural shape. But the indexer's scoring math is -`ReLU(q·k) · w_h` summed across heads — TurboQuant's "preserve inner -products" guarantee from Algorithm 2 doesn't compose with the ReLU -nonlinearity. The quantization error becomes worst-case at the threshold, -which is where top-k decisions get made. **Bad fit precisely where it -matters most.** - -### 3. NVFP4 hardware exists; TurboQuant is software-only - -TurboQuant runs as bit-packed uint8 + Triton kernels. It can't use -tcgen05 FP4 tensor cores because its values aren't FP4 — they're -codebook *indices*. So you'd be paying CPU/GPU cycles to dequant via -gathers and per-token rotation matrix-vector multiplies, when the same -storage cost (4 bits/value) is available natively as FP4 with hardware -dequant during MMA. - -The TurboQuant benchmark numbers (+3–5% throughput at 3-bit) are -real, but they're against `bf16_kv` baselines on architectures that -don't have FP4 tensor cores. On Blackwell with NVFP4, the comparison -should be FP4 storage + FP4 MMA — which is strictly better in every -axis (bandwidth, capacity, dequant cost). - -## Where TurboQuant *would* help, and the verdict on whether it's worth it - -The only DSv4 stream where TurboQuant's shape is a natural fit is the -**SWA branch** — uncompressed per-token KV in the sliding window, 128 -tokens × `n_layers` × `hd=512` = 8 MB at 1M context. - -**It's 8 MB.** Not worth a new dependency, a paper-grade extra failure -mode, or the rotation overhead. The SWA branch fits in L2 cache on B200. - -### Verdict - -Don't use TurboQuant. The right move for DSv4's KV cache is **FP4 storage -+ FP4 MMA on the compressed streams**, fully Blackwell-native, paper- -aligned (§5.2.1), with no codebook lookup overhead. The infrastructure to -do this is already in your kernel library (the `ENABLE_FP4_EPILOGUE` -template, the FP4 MMA path). - -If you want a paper to cite for "what's the state-of-the-art KV -compression in 2026," TurboQuant is one. If you want the highest-perf -production-grade DSv4 implementation, native FP4 is the answer. - ---- # PRIORITY ORDER (updated 2026-06-02) @@ -312,18 +221,13 @@ production-grade DSv4 implementation, native FP4 is the answer. | **P1** | Same for shared expert | S | ~120 launches/token | ✅ Done | | **P2** | Drop per-call `fill_()` in Nvfp4Linear | S | ~244 launches/token | ✅ Done | | **P3** | CUDA RoPE kernel (1 launch vs 5-6) | S | ~732 launches/token | ✅ Done | -| **KV-1** | FP4 storage for CSA main compressed KV | M | Huge at long context | Next | -| **KV-2** | FP4 storage for HCA compressed KV | M | Same pattern as KV-1 | After KV-1 | -| **KV-3** | FP4 storage for indexer keys (pair with E7) | M | Throughput + paper compliance | After KV-2 | +| **KV-1** | FP4 storage for CSA main compressed KV | M | Huge at long context | Next | ✅ Done | +| **KV-2** | FP4 storage for HCA compressed KV | M | Same pattern as KV-1 | After KV-1 | ✅ Done | +| **KV-3** | FP4 storage for indexer keys (pair with E7) | M | Throughput + paper compliance | After KV-2 |✅ Done | | **P4** | RMSNorm fused into next quantize | S | 122 launches/token | After KV | | **P5** | mHC pre_block + RMSNorm fused | S | ~120 launches/token | After P4 | | **P6** | CUDA graph capture | L | **2–3× total** | After everything above | -**Part 1 complete.** The NVFP4-everywhere gap for the GEMM+activation+RoPE -path is closed. The remaining wins are KV-cache dtype (Part 2) and -higher-order fusion (P4–P6). Land all of those before attempting CUDA -graphs — the captured graph should reflect the final fused structure, not -the pre-fusion one. --- @@ -354,9 +258,3 @@ the pre-fusion one. ms**, not three-digit. If it isn't, something is still on the hot path that shouldn't be, and the answer is "profile, don't guess next." - -6. **Don't optimize for problems you don't have.** TurboQuant is the - cautionary tale here. The KV cache at 1M is 10 GB on 8 × B200 — that - is not a problem that needs solving with a new dependency. The - problem is throughput, and the right answer is FP4 storage + FP4 MMA, - which is hardware-native and doesn't require codebook lookups. diff --git a/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu index bdd71bc9..2af431b5 100644 --- a/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu +++ b/dsv4/kernels/cuda/fused_rmsnorm_quantize.cu @@ -1,39 +1,60 @@ -/* - * fused_rmsnorm_amax_quantize.cu +/** + * fused_rmsnorm_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 + * Replaces: rmsnorm (4+ BF16 launches) + amax (1 launch) + quantize (1 launch) + * with just 2 kernel launches. * - * 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 + * Kernel 1: rmsnorm_amax_gsa_kernel + * - Compute RMS of each row: rms = sqrt(mean(x^2) + eps) + * - Compute row-wise amax of (x / rms * weight) — the normalized output + * - Derive gsa = amax / divisor for each row + * - Write gsa (per-row) and inv_rms (per-row) to GPU buffers * - * Production: used for attention input norm, FFN input norm, q_a_norm, kv_norm. + * 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 + * - 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 */ #include +#include #include #include +#include +#include +#include +#include #include #include -// FP4 E2M1 look-up table (same as production) +// 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}; // ============================================================================ -// Kernel 1: Compute RMS + row-wise amax + gsa +// Kernel 1: Compute RMS + amax of normalized output → gsa per row // ============================================================================ -// Each block handles one row. Threadblock: (WARP_SIZE, 1, 1) per row. -// For hidden_size=7168: 7168/32 = 224 warps needed. Use multiple blocks. +// 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 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) + float* __restrict__ inv_rms_out, // (M,) FP32 — per-row 1/rms (for kernel 2) const int M, const int N, const float eps, @@ -41,201 +62,280 @@ __global__ void rmsnorm_amax_gsa_kernel( ) { 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 + + 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 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; + sum_sq += val * val; + } + + // Warp-level reduction + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + sum_sq += __shfl_down_sync(0xFFFFFFFF, sum_sq, offset); + } + + // Block-level reduction via shared memory + const int num_warps = blockDim.x / warpSize; + __shared__ float s_sum_sq[32]; // max 32 warps + int lane = threadIdx.x % warpSize; + int warp_id = threadIdx.x / warpSize; + + if (lane == 0) { + s_sum_sq[warp_id] = sum_sq; + } + __syncthreads(); + + // First warp reduces across warps + float row_sum_sq = 0.0f; + if (warp_id == 0) { + row_sum_sq = (lane < num_warps) ? s_sum_sq[lane] : 0.0f; + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + row_sum_sq += __shfl_down_sync(0xFFFFFFFF, row_sum_sq, offset); + } + } + + // Broadcast inv_rms to all threads + __shared__ float s_inv_rms; + if (threadIdx.x == 0) { + float rms = sqrtf(row_sum_sq / N + eps); + s_inv_rms = 1.0f / rms; + } + __syncthreads(); + float inv_rms = s_inv_rms; + + // Sub-pass 2: amax of normalized output (x * inv_rms * weight) + float row_amax = 0.0f; + for (int col = threadIdx.x; col < N; col += blockDim.x) { + float val = __bfloat162float(x_row[col]) * inv_rms * norm_weight[col]; float abs_val = fabsf(val); if (abs_val > row_amax) row_amax = abs_val; } - - // Warp-level reduction for sum_sq and amax + + // Warp-level reduce max 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; + float global_amax = 0.0f; + if (lane < num_warps) global_amax = s_amax[lane]; 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)); + global_amax = fmaxf(global_amax, __shfl_down_sync(0xFFFFFFFF, global_amax, offset)); + } + if (lane == 0) { + gsa_out[row] = fmaxf(global_amax, 1e-8f) / divisor; + inv_rms_out[row] = inv_rms; } - } - - 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. +// 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; +} __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 + const float* __restrict__ inv_rms, // (M,) FP32 — per-row 1/rms uint8_t* __restrict__ x_fp4, // (M, N//2) FP4 packed - __nv_fp8_e4m3* __restrict__ x_sf, // (M, N//16) E4M3 block scales + uint8_t* __restrict__ x_sf, // (M, N//16) E4M3 block scales const int M, - const int N, - const float eps + const int N ) { - const int row = blockIdx.x; + const int row = blockIdx.y; + const int n_block = blockIdx.x; if (row >= M) return; - - const __nv_bfloat16* x_row = x + row * N; + if (n_block * 16 >= N) return; + + const __nv_bfloat16* x_row = x + (size_t)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 + float row_inv_rms = inv_rms[row]; + + // Step 1: Load 16 BF16 elements, normalize, compute block amax + float vals[16]; + float block_amax = 0.0f; + const int col_base = n_block * 16; + + for (int i = 0; i < 16; i++) { + int col = col_base + i; + if (col < N) { + float v = __bfloat162float(x_row[col]); + v = v * row_inv_rms * norm_weight[col]; // RMSNorm vals[i] = v; float av = fabsf(v); if (av > block_amax) block_amax = av; + } else { + vals[i] = 0.0f; } - - // 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; } + + // 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); + + // 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 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; + for (int i = 0; i < 8; i++) { + dst_fp4[i] = fp4_bytes[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; +} + +// ============================================================================ +// PyTorch bridge +// ============================================================================ + +std::tuple +rmsnorm_quantize_nvfp4_cuda( + torch::Tensor x, // (M, N) BF16 + torch::Tensor norm_weight, // (N,) FP32 + double eps, + double divisor +) { + TORCH_CHECK(x.is_contiguous(), "x must be contiguous"); + TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be BF16"); + TORCH_CHECK(norm_weight.scalar_type() == torch::kFloat32, "norm_weight must be FP32"); + + const int M = x.size(0); + const int N = x.size(1); + TORCH_CHECK(N % 16 == 0, "N must be multiple of 16"); + + auto stream = c10::cuda::getCurrentCUDAStream(); + auto options = x.options(); + + // Output buffers + 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)); + + // Kernel 1: RMSNorm + amax → gsa (1 block per row) + const int threads1 = 256; // 8 warps, handles up to N=8192 + rmsnorm_amax_gsa_kernel<<>>( + reinterpret_cast(x.data_ptr()), + norm_weight.data_ptr(), + gsa.data_ptr(), + inv_rms.data_ptr(), + M, N, (float)eps, (float)divisor + ); + + // Kernel 2: Normalize + quantize (1 block per (row, microblock)) + const int n_blocks = N / 16; + dim3 grid2(n_blocks, M); + const int threads2 = 16; // 1 thread per element in the 16-elem microblock + rmsnorm_quantize_nvfp4_kernel<<>>( + reinterpret_cast(x.data_ptr()), + norm_weight.data_ptr(), + gsa.data_ptr(), + inv_rms.data_ptr(), + x_fp4.data_ptr(), + reinterpret_cast(x_sf.data_ptr()), + M, N + ); + + return std::make_tuple(x_fp4, x_sf, gsa, inv_rms); +} + +// Standalone kernel 1 entry point (for testing / when only gsa needed) +torch::Tensor rmsnorm_amax_gsa_cuda( + torch::Tensor x, + torch::Tensor norm_weight, + double eps, + 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"); + + const int M = x.size(0); + const int N = x.size(1); + auto stream = c10::cuda::getCurrentCUDAStream(); + + auto gsa = torch::empty({M}, x.options().dtype(torch::kFloat32)); + auto inv_rms = torch::empty({M}, x.options().dtype(torch::kFloat32)); + + const int threads = 256; + rmsnorm_amax_gsa_kernel<<>>( + reinterpret_cast(x.data_ptr()), + norm_weight.data_ptr(), + gsa.data_ptr(), + inv_rms.data_ptr(), + M, N, (float)eps, (float)divisor + ); + + return gsa; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("rmsnorm_quantize_nvfp4", &rmsnorm_quantize_nvfp4_cuda, + "Fused RMSNorm + amax + quantize to NVFP4"); + m.def("rmsnorm_amax_gsa", &rmsnorm_amax_gsa_cuda, + "RMSNorm + amax → gsa (kernel 1 only)"); } diff --git a/dsv4/ops/quantize.py b/dsv4/ops/quantize.py index 72b7f2ae..8342762d 100644 --- a/dsv4/ops/quantize.py +++ b/dsv4/ops/quantize.py @@ -349,3 +349,54 @@ def quantize_nvfp4_gpu(x_bf16, global_scale): from dsv4.kernels.cuda.loader import get_cuda_module mod = get_cuda_module("quantize_nvfp4", ["quantize_nvfp4.cu"]) return mod.quantize_nvfp4(x_bf16, global_scale) + + +def dequantize_nvfp4(x_fp4, x_sf, gsa, shape=None): + """Dequantize NVFP4 → BF16 using the CUDA dequant kernel. + + Args: + x_fp4: (M, N//2) FP4 packed + x_sf: (M, N//16) E4M3 block scales + gsa: (M,) or (M, 1) or (1,) FP32 global scale per row + shape: unused, kept for API compat + + Returns: + (M, N) BF16 tensor + """ + from dsv4.kernels.cuda.loader import get_cuda_module + mod = get_cuda_module("dequant_nvfp4", ["dequant_nvfp4.cu"]) + if gsa.dim() == 2: + gsa = gsa.squeeze(1) # (M, 1) → (M,) + return mod.dequant_nvfp4(x_fp4, x_sf, gsa) + + +def rmsnorm_quantize_nvfp4(x_bf16, norm_weight, eps=1e-6, divisor=6.0 * 448.0): + """Fused RMSNorm + amax + NVFP4 quantize: 2 kernel launches total. + + Replaces the unfused path: + rmsnorm(x, weight) → 4+ BF16 launches + quantize_nvfp4_gpu_fused(rmsnormed) → 2 kernel launches + amax + Total unfused: 6+ launches per call × 122 calls/layer-step = 732+ launches/token + + Fused: 2 kernel launches per call × 122 calls = 244 launches → 488 launches saved/token. + + Two-kernel approach (correct cross-CTA reduction): + Kernel 1: compute RMS + amax of normalized output → gsa per row (GPU buffer) + Kernel 2: normalize + quantize using gsa from GPU buffer (no CPU sync) + + Args: + x_bf16: (M, N) BF16 tensor. N must be a multiple of 16. + norm_weight: (N,) FP32 RMSNorm weight. + eps: RMSNorm epsilon (default 1e-6). + divisor: gsa = amax / divisor. Default 6.0 * 448.0 = 2688.0. + + Returns: + x_fp4: (M, N//2) FP4 packed (uint8 view of float4_e2m1fn_x2) + x_sf: (M, N//16) E4M3 block scales + gsa: (M,) FP32 per-row global scale for GEMM + inv_rms: (M,) FP32 per-row 1/RMS (useful for downstream if needed) + """ + from dsv4.kernels.cuda.loader import get_cuda_module + mod = get_cuda_module("fused_rmsnorm_quantize", ["fused_rmsnorm_quantize.cu"]) + x_fp4, x_sf, gsa, inv_rms = mod.rmsnorm_quantize_nvfp4(x_bf16, norm_weight, eps, divisor) + return x_fp4, x_sf, gsa, inv_rms diff --git a/tests/unit/test_fused_rmsnorm_quantize.py b/tests/unit/test_fused_rmsnorm_quantize.py new file mode 100644 index 00000000..4633d3e5 --- /dev/null +++ b/tests/unit/test_fused_rmsnorm_quantize.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Test fused RMSNorm + NVFP4 quantize kernel. + +Validates the fused kernel against the reference (unfused) path: + Reference: rmsnorm(x, weight) → quantize_nvfp4_gpu_fused(rmsnormed) + Fused: rmsnorm_quantize_nvfp4(x, weight, eps, divisor) → (fp4, sf, gsa, inv_rms) + +Tests at production shapes: (1, 7168) for decode, (8, 7168) for prefill. +""" + +import torch +import math + +def rmsnorm_ref(x, weight, eps=1e-6): + """PyTorch reference RMSNorm.""" + xf = x.float() + rms = xf.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt() + return xf * rms * weight.float() + + +def main(): + from dsv4.ops.quantize import quantize_nvfp4_gpu_fused, rmsnorm_quantize_nvfp4, dequantize_nvfp4 + + device = "cuda" + torch.manual_seed(42) + + DIVISOR = 6.0 * 448.0 # same as production + EPS = 1e-6 + HIDDEN = 7168 # production hidden_size + + test_configs = [ + (1, HIDDEN, "decode T=1"), + (8, HIDDEN, "prefill T=8"), + (128, HIDDEN, "prefill T=128"), + ] + + all_pass = True + for M, N, label in test_configs: + print(f"\n=== Test: {label} (M={M}, N={N}) ===") + + x = torch.randn(M, N, dtype=torch.bfloat16, device=device) + weight = torch.randn(N, dtype=torch.float32, device=device) + # Ensure weight is positive-ish (norm weights are near 1.0 in practice) + weight = weight * 0.1 + 1.0 + + # Reference path: unfused + x_normed = rmsnorm_ref(x, weight, EPS) # (M, N) FP32 + x_normed_bf16 = x_normed.bfloat16() + + # Unfused quantize + x_fp4_ref, x_sf_ref, gsa_ref = quantize_nvfp4_gpu_fused(x_normed_bf16, DIVISOR) + + # Fused path + x_fp4_fused, x_sf_fused, gsa_fused, inv_rms_fused = rmsnorm_quantize_nvfp4( + x, weight, EPS, DIVISOR + ) + + # Validate shapes + assert x_fp4_fused.shape == (M, N // 2), f"fp4 shape: {x_fp4_fused.shape} != {(M, N//2)}" + assert x_sf_fused.shape == (M, N // 16), f"sf shape: {x_sf_fused.shape} != {(M, N//16)}" + assert gsa_fused.shape == (M,), f"gsa shape: {gsa_fused.shape} != {(M,)}" + assert inv_rms_fused.shape == (M,), f"inv_rms shape: {inv_rms_fused.shape} != {(M,)}" + + # Validate gsa matches + # gsa is per-row; unfused computes a scalar from the whole tensor + # For T=1 they should match exactly. For T>1, fused is per-row, unfused is scalar. + if M == 1: + gsa_diff = (gsa_fused[0] - gsa_ref[0]).abs().item() + print(f" gsa: fused={gsa_fused[0].item():.6f} ref={gsa_ref[0].item():.6f} diff={gsa_diff:.2e}") + if gsa_diff > 1e-3: + print(f" FAIL: gsa mismatch") + all_pass = False + + # Dequantize and compare end-to-end + dq_fused = dequantize_nvfp4(x_fp4_fused, x_sf_fused, gsa_fused) + dq_ref = dequantize_nvfp4(x_fp4_ref, x_sf_ref, gsa_ref) + + # Compute cosine similarity of dequantized outputs + cos = torch.nn.functional.cosine_similarity( + dq_fused.float().flatten().unsqueeze(0), + dq_ref.float().flatten().unsqueeze(0) + ).item() + print(f" Dequant cosine (fused vs unfused): {cos:.6f}") + + # Also compare against the true RMSNorm output + cos_vs_ref = torch.nn.functional.cosine_similarity( + dq_fused.float().flatten().unsqueeze(0), + x_normed.flatten().unsqueeze(0) + ).item() + print(f" vs true RMSNorm: {cos_vs_ref:.6f}") + + if cos < 0.998: + print(f" FAIL: dequant cosine too low ({cos:.6f})") + all_pass = False + if cos_vs_ref < 0.990: + print(f" FAIL: vs true RMSNorm cosine too low ({cos_vs_ref:.6f})") + all_pass = False + + # Test: unweighted RMSNorm (for q_b norm and kv norm) + print("\n=== Test: unweighted RMSNorm (weight=1.0) ===") + for M, N, label in [(1, HIDDEN, "decode"), (8, HIDDEN, "prefill")]: + x = torch.randn(M, N, dtype=torch.bfloat16, device=device) + weight_ones = torch.ones(N, dtype=torch.float32, device=device) + + x_normed = rmsnorm_ref(x, weight_ones, EPS).bfloat16() + x_fp4_ref, x_sf_ref, gsa_ref = quantize_nvfp4_gpu_fused(x_normed, DIVISOR) + + x_fp4_fused, x_sf_fused, gsa_fused, inv_rms_fused = rmsnorm_quantize_nvfp4( + x, weight_ones, EPS, DIVISOR + ) + + dq_fused = dequantize_nvfp4(x_fp4_fused, x_sf_fused, gsa_fused) + dq_ref = dequantize_nvfp4(x_fp4_ref, x_sf_ref, gsa_ref) + + cos = torch.nn.functional.cosine_similarity( + dq_fused.float().flatten().unsqueeze(0), + dq_ref.float().flatten().unsqueeze(0) + ).item() + print(f" {label}: cos={cos:.6f}") + if cos < 0.998: + print(f" FAIL") + all_pass = False + + # Test: validate gsa against per-row reference + print("\n=== Test: gsa per-row validation ===") + for M, N, label in [(1, HIDDEN, "decode"), (8, HIDDEN, "prefill")]: + x = torch.randn(M, N, dtype=torch.bfloat16, device=device) + weight = torch.randn(N, dtype=torch.float32, device=device) * 0.1 + 1.0 + + x_fp4_fused, x_sf_fused, gsa_fused, inv_rms_fused = rmsnorm_quantize_nvfp4( + x, weight, EPS, DIVISOR + ) + + max_diff = 0.0 + for row in range(M): + x_normed_row = rmsnorm_ref(x[row:row+1], weight, EPS) + ref_amax = x_normed_row.abs().max().item() + ref_gsa = max(ref_amax, 1e-8) / DIVISOR + diff = abs(gsa_fused[row].item() - ref_gsa) + max_diff = max(max_diff, diff) + + print(f" {label}: max gsa diff vs reference = {max_diff:.2e}") + if max_diff > 0.1: # FP4 quantization + BF16 round-trip introduces some error + print(f" FAIL: gsa diff too large") + all_pass = False + + print(f"\n{'='*60}") + print(f"{'ALL TESTS PASSED' if all_pass else 'SOME TESTS FAILED'}") + return 0 if all_pass else 1 + + +if __name__ == "__main__": + exit(main())