P4: Fused RMSNorm + NVFP4 quantize kernel (2 launches vs 6+)
- fused_rmsnorm_quantize.cu: two-kernel approach Kernel 1: rmsnorm_amax_gsa — compute RMS + amax of normalized output → gsa per row Kernel 2: rmsnorm_quantize_nvfp4 — normalize + quantize using GPU-computed gsa - Python bridge: rmsnorm_quantize_nvfp4() in ops/quantize.py - Python bridge: dequantize_nvfp4() in ops/quantize.py - Unit test: test_fused_rmsnorm_quantize.py (production shapes: 7168 hidden) - Eliminates ~488 kernel launches per token (122 sites × 4 launches saved)
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <ATen/ATen.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/extension.h>
|
||||
#include <cstdint>
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
|
||||
// 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<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
|
||||
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<<<M, threads1, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr<at::BFloat16>()),
|
||||
norm_weight.data_ptr<float>(),
|
||||
gsa.data_ptr<float>(),
|
||||
inv_rms.data_ptr<float>(),
|
||||
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<<<grid2, threads2, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr<at::BFloat16>()),
|
||||
norm_weight.data_ptr<float>(),
|
||||
gsa.data_ptr<float>(),
|
||||
inv_rms.data_ptr<float>(),
|
||||
x_fp4.data_ptr<uint8_t>(),
|
||||
reinterpret_cast<uint8_t*>(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<<<M, threads, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr<at::BFloat16>()),
|
||||
norm_weight.data_ptr<float>(),
|
||||
gsa.data_ptr<float>(),
|
||||
inv_rms.data_ptr<float>(),
|
||||
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)");
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
153
tests/unit/test_fused_rmsnorm_quantize.py
Normal file
153
tests/unit/test_fused_rmsnorm_quantize.py
Normal file
@@ -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())
|
||||
Reference in New Issue
Block a user