From cacf64232e7b82694189370f3d0372134a855dab Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 21:26:51 +0000 Subject: [PATCH] CRITICAL FIX: fused_amax_quantize cross-CTA race condition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The single-kernel approach used __syncthreads() for cross-CTA amax reduction, but __syncthreads() only syncs within a CTA (same blockIdx). CTA 0 reading s_amax[1] before CTA 1 writes = race condition = garbage gsa. Result: residual |X| exploded to 10^37 by L0. F_attn and F_ffn were 0.0. Fix: Two-kernel approach (correct, zero CPU syncs): Kernel 1: amax_gsa.cu — computes gsa on GPU, returns GPU tensor Kernel 2: quantize_nvfp4_from_buffer — reads gsa from GPU buffer The fused_amax_quantize.cu now exports quantize_nvfp4_from_buffer and deinterleave_quantize_from_buffer (gsa from GPU buffer, not kernel param). Same P0 win: zero .item() syncs. Two kernel launches instead of one, but correctness > shaving one launch. --- dsv4/kernels/cuda/fused_amax_quantize.cu | 220 ++++++++++++++--------- dsv4/kernels/cuda/loader.py | 6 +- dsv4/ops/quantize.py | 43 +++-- 3 files changed, 167 insertions(+), 102 deletions(-) diff --git a/dsv4/kernels/cuda/fused_amax_quantize.cu b/dsv4/kernels/cuda/fused_amax_quantize.cu index 6ddb3006..ca143183 100644 --- a/dsv4/kernels/cuda/fused_amax_quantize.cu +++ b/dsv4/kernels/cuda/fused_amax_quantize.cu @@ -1,35 +1,43 @@ /** * Fused amax + gsa + NVFP4 quantization kernel. * - * Single kernel launch that: - * 1. Computes row-wise amax of the input (GPU-only, no CPU sync) - * 2. Derives gsa = max(amax) / divisor - * 3. Quantizes each row to NVFP4 (FP4 data + FP8 E4M3 block scales) - * 4. Writes gsa to a GPU buffer for downstream GEMM global_scale_a - * - * This eliminates ALL .item() syncs from the NVFP4 activation path. - * Previously: quantize_nvfp4.cu + amax_gsa.cu required: - * - .item() for amax (amax_gsa path) - * - .item() to pass gsa as kernel param (quantize_nvfp4 path) - * Now: zero CPU-GPU syncs. gsa stays on GPU. - * - * Grid: (N / 16, M, 1) — each CTA processes one 16-element block in one row. - * Block: 256 threads (for cross-CTA amax reduction in the y-dimension). - * - * The amax reduction uses a two-phase approach: - * Phase 1: Each CTA computes its local max |x| (across the 16 elements it quantizes) - * Phase 2: The CTA at n_block=0 reduces across all n_blocks via shared memory + * Two-phase approach: + * Phase 1: Each CTA quantizes its 16-element block (independent). + * Phase 2: CTA 0 of each row reduces across all CTAs via atomicMax * to get the row-wide amax, then derives gsa. * - * For decode (M=1, N=7168, 448 CTAs): amax is computed in the same kernel - * as quantization. No separate kernel launch, no CPU sync. + * The amax reduction uses global memory atomics (not shared memory) + * to correctly handle cross-CTA synchronization within the same kernel. + * Each CTA writes its block_amax to a global memory buffer. + * After a grid-sync (via cooperative groups or a second launch), + * CTA 0 computes the row-wide amax from all block amaxes. * - * For batched decode (M>1): each row is independent. gsa is per-row. - * The GEMM path uses a single gsa (max across all rows), which we compute - * by having the first CTA of each row write its row's amax to a buffer, - * then a final single-CTA pass reduces across rows. - * - * For single-token decode (M=1), the final pass is trivially just one value. + * Since we can't do a proper grid sync in a single kernel without + * cooperative groups (which requires special launch), we use a two-kernel + * approach instead: + * Kernel 1: Compute per-block amaxes + quantize to NVFP4. + * Kernel 2: Reduce per-block amaxes to per-row gsa. + * + * Actually, the simplest correct approach is: + * - Compute gsa in a separate lightweight kernel (amax_gsa.cu already does this) + * - Pass gsa as a GPU buffer to quantize_nvfp4 + * - quantize_nvfp4 reads gsa from the GPU buffer instead of a kernel param + * + * This file implements the SINGLE-CTA-per-row case (N <= 16). + * For the general case, use the two-kernel approach. + * + * UPDATE: Switched to per-CTA-independent quantize with a global amax + * reduction. Each CTA computes its own amax, writes to a global buffer. + * A final pass (CTA 0 per row) reads all amaxes and computes gsa. + * But this requires grid sync which we don't have. + * + * SIMPLEST CORRECT APPROACH: + * Use the existing amax_gsa.cu kernel to compute gsa on GPU, + * then pass the GPU tensor to quantize_nvfp4 via a modified kernel + * that reads global_scale from a GPU buffer instead of a kernel parameter. + * + * This file is KEPT but the quantize kernel is modified to accept + * global_scale from a GPU buffer. */ #include @@ -40,7 +48,6 @@ #include #include #include -#include __device__ __forceinline__ int half_step_to_e2m1(int hs) { if (hs <= 4) return hs; @@ -50,78 +57,101 @@ __device__ __forceinline__ int half_step_to_e2m1(int hs) { return 7; } -// Shared memory layout: one float per n_block for amax reduction -// Max N/16 CTAs per row. For N=7168, that's 448. 448 floats = 1792 bytes. -// We use dynamic shared memory. - -__global__ void fused_amax_quantize_nvfp4_kernel( +/** + * Quantize kernel that reads global_scale from a GPU buffer. + * Same as quantize_nvfp4.cu but gsa comes from GMEM, not a kernel param. + * This enables zero-CPU-sync operation: gsa computed on GPU → passed directly. + */ +__global__ void quantize_nvfp4_from_buffer_kernel( const __nv_bfloat16* __restrict__ input, int M, int N, - float divisor, + const float* __restrict__ gsa_buffer, // (M,) GPU buffer with per-row gsa uint8_t* __restrict__ out_fp4, - uint8_t* __restrict__ out_sf, - float* __restrict__ out_gsa // (M,) GPU buffer — gsa per row + uint8_t* __restrict__ out_sf ) { int m = blockIdx.y; int n_block = blockIdx.x; - int n_blocks = gridDim.x; if (m >= M || n_block * 16 >= N) return; - extern __shared__ float s_amax[]; + float gsa = gsa_buffer[m]; - // Step 1: Read 16 BF16 elements and compute local amax float vals[16]; float block_amax = 0.0f; + + // Step 1: Read 16 BF16 elements and compute amax for (int i = 0; i < 16; i++) { int col = n_block * 16 + i; if (col < N) { - vals[i] = __bfloat162float(input[m * N + col]); + vals[i] = __bfloat162float(input[m * N + col]) / gsa; } else { vals[i] = 0; } block_amax = fmaxf(block_amax, fabsf(vals[i])); } - // Step 2: Cross-CTA reduction to get row-wide amax - // Each CTA writes its local amax to shared memory, then CTA 0 reduces. - if (n_block < n_blocks) { - s_amax[n_block] = block_amax; + // Step 2: Compute FP8 E4M3 block scale + float bsf = block_amax / 6.0f; + if (block_amax < 6.0f * 0.001953125f) { + bsf = 0; + for (int i = 0; i < 16; i++) vals[i] = 0; } - __syncthreads(); + __nv_fp8_e4m3 bsf8_obj(bsf); + float bs = (float)bsf8_obj; + uint8_t bsf8 = *(uint8_t*)&bsf8_obj; - // CTA 0 computes the row-wide amax and derives gsa - float gsa; - if (n_block == 0) { - float row_amax = 0.0f; - for (int b = 0; b < n_blocks; b++) { - row_amax = fmaxf(row_amax, s_amax[b]); - } - gsa = fmaxf(row_amax, 1e-8f) / divisor; - out_gsa[m] = gsa; // Write gsa to GPU buffer for GEMM - } - // Broadcast gsa from CTA 0 to all CTAs via shared memory - __syncthreads(); - // Re-read from the s_amax[0] slot where CTA 0 stored gsa temporarily - // Actually, we need a different approach since s_amax is being used. - // Store gsa in a known location. - if (n_block == 0) { - s_amax[0] = gsa; - } - __syncthreads(); - gsa = s_amax[0]; - - // Step 3: Quantize — divide by gsa, compute FP8 block scale, quantize to FP4 + // Step 3: Quantize each value to FP4 E2M1 + uint8_t nibbles[16]; for (int i = 0; i < 16; i++) { - vals[i] = vals[i] / gsa; + if (bs < 1e-8f) { nibbles[i] = 0; continue; } + float s = vals[i] / bs; + 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; } - float q_amax = 0.0f; + // Step 4: Pack pairs + for (int i = 0; i < 8; i++) + out_fp4[m * (N / 2) + n_block * 8 + i] = (nibbles[2*i+1] << 4) | nibbles[2*i]; + + // Step 5: Write FP8 block scale + out_sf[m * (N / 16) + n_block] = bsf8; +} + +/** + * Deinterleave + quantize kernel that reads global_scale from a GPU buffer. + * For the MoE fused_swiglu L2 path. + */ +__global__ void deinterleave_quantize_from_buffer_kernel( + const __nv_bfloat16* __restrict__ fused, + int M, int N, int intermediate, int granularity, + const float* __restrict__ gsa_buffer, + uint8_t* __restrict__ out_fp4, + uint8_t* __restrict__ out_sf +) { + int m = blockIdx.y; + int n_block = blockIdx.x; + if (m >= M || n_block * 16 >= intermediate) return; + + float gsa = gsa_buffer[m]; + + float vals[16]; + float block_amax = 0.0f; + for (int i = 0; i < 16; i++) { - q_amax = fmaxf(q_amax, fabsf(vals[i])); + int nd = n_block * 16 + i; + if (nd >= intermediate) { vals[i] = 0; continue; } + int group = 2 * (nd / granularity) + 1; + int offset = nd % granularity; + int fc = group * granularity + offset; + float v = __bfloat162float(fused[m * N + fc]); + vals[i] = v / gsa; + block_amax = fmaxf(block_amax, fabsf(vals[i])); } - float bsf = q_amax / 6.0f; - if (q_amax < 6.0f * 0.001953125f) { + float bsf = block_amax / 6.0f; + if (block_amax < 6.0f * 0.001953125f) { bsf = 0; for (int i = 0; i < 16; i++) vals[i] = 0; } @@ -141,36 +171,54 @@ __global__ void fused_amax_quantize_nvfp4_kernel( } for (int i = 0; i < 8; i++) - out_fp4[m * (N / 2) + n_block * 8 + i] = (nibbles[2*i+1] << 4) | nibbles[2*i]; + out_fp4[m * (intermediate / 2) + n_block * 8 + i] = (nibbles[2*i+1] << 4) | nibbles[2*i]; - out_sf[m * (N / 16) + n_block] = bsf8; + out_sf[m * (intermediate / 16) + n_block] = bsf8; } -std::tuple fused_amax_quantize_nvfp4_cuda( - torch::Tensor input_bf16, double divisor +// Python API: quantize with gsa from GPU buffer +std::tuple quantize_nvfp4_from_buffer_cuda( + torch::Tensor input_bf16, torch::Tensor gsa_buffer ) { int M = input_bf16.size(0); int N = input_bf16.size(1); - TORCH_CHECK(N % 16 == 0, "N must be a multiple of 16 for NVFP4 quantization"); + TORCH_CHECK(N % 16 == 0, "N must be a multiple of 16"); + TORCH_CHECK(gsa_buffer.size(0) == M, "gsa_buffer size must match M"); auto opts = input_bf16.options(); auto out_fp4 = torch::zeros({M, N / 2}, opts.dtype(torch::kUInt8)); auto out_sf = torch::zeros({M, N / 16}, opts.dtype(torch::kUInt8)); - auto out_gsa = torch::zeros({M}, opts.dtype(torch::kFloat32)); - int nb = N / 16; dim3 grid(nb, M); dim3 block(16); - int smem_size = nb * sizeof(float); - - fused_amax_quantize_nvfp4_kernel<<>>( + quantize_nvfp4_from_buffer_kernel<<>>( reinterpret_cast(input_bf16.data_ptr()), - M, N, (float)divisor, - out_fp4.data_ptr(), out_sf.data_ptr(), - out_gsa.data_ptr() + M, N, gsa_buffer.data_ptr(), + out_fp4.data_ptr(), out_sf.data_ptr() ); - return {out_fp4.view(torch::kFloat4_e2m1fn_x2), out_sf.view(torch::kFloat8_e4m3fn), out_gsa}; + return {out_fp4.view(torch::kFloat4_e2m1fn_x2), out_sf.view(torch::kFloat8_e4m3fn)}; +} + +// Python API: deinterleave + quantize with gsa from GPU buffer +std::tuple deinterleave_quantize_from_buffer_cuda( + torch::Tensor fused_bf16, int64_t intermediate, int64_t granularity, torch::Tensor gsa_buffer +) { + int M = fused_bf16.size(0); + int N = fused_bf16.size(1); + auto opts = fused_bf16.options(); + auto out_fp4 = torch::zeros({M, (int)intermediate / 2}, opts.dtype(torch::kUInt8)); + auto out_sf = torch::zeros({M, (int)intermediate / 16}, opts.dtype(torch::kUInt8)); + int nb = (int)intermediate / 16; + dim3 grid(nb, M); + dim3 block(16); + deinterleave_quantize_from_buffer_kernel<<>>( + reinterpret_cast(fused_bf16.data_ptr()), + M, N, (int)intermediate, (int)granularity, gsa_buffer.data_ptr(), + out_fp4.data_ptr(), out_sf.data_ptr() + ); + return {out_fp4.view(torch::kFloat4_e2m1fn_x2), out_sf.view(torch::kFloat8_e4m3fn)}; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("fused_amax_quantize_nvfp4", &fused_amax_quantize_nvfp4_cuda); + m.def("quantize_nvfp4_from_buffer", &quantize_nvfp4_from_buffer_cuda); + m.def("deinterleave_quantize_from_buffer", &deinterleave_quantize_from_buffer_cuda); } diff --git a/dsv4/kernels/cuda/loader.py b/dsv4/kernels/cuda/loader.py index d4f9bdf9..2f168f8c 100644 --- a/dsv4/kernels/cuda/loader.py +++ b/dsv4/kernels/cuda/loader.py @@ -67,9 +67,11 @@ def get_cuda_module(name, sources, extra_cuda_cflags=None): def preload_all(): """Preload all CUDA kernels at startup (before the hot path).""" - # Fused amax + quantize — THE critical kernel for P0 + # amax_gsa — computes gsa on GPU (no .item()) + get_cuda_module("amax_gsa", ["amax_gsa.cu"]) + # quantize-from-buffer — reads gsa from GPU buffer (no .item()) get_cuda_module("fused_amax_quantize", ["fused_amax_quantize.cu"]) - # Standalone quantize (used by weight quantization, not hot path) + # Standalone quantize (for when gsa is known, not hot path) get_cuda_module("quantize_nvfp4", ["quantize_nvfp4.cu"]) # Sampler get_cuda_module("sampler", ["sampler.cu"]) diff --git a/dsv4/ops/quantize.py b/dsv4/ops/quantize.py index 2777a73f..5ebef74d 100644 --- a/dsv4/ops/quantize.py +++ b/dsv4/ops/quantize.py @@ -248,10 +248,11 @@ def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, gra def deinterleave_amax_quantize_nvfp4_fused(fused_bf16, intermediate, divisor=6.0 * 448.0, granularity=8): - """Fused deinterleave + amax + gsa + quantize: NO CPU sync, single kernel launch. + """Fused deinterleave + amax + quantize: zero CPU syncs, two kernel launches. - For the MoE fused_swiglu L2 path. Computes gsa from the de-interleaved - (SwiGLU) values on GPU, quantizes in the same kernel. Zero .item() syncs. + For the MoE fused_swiglu L2 path. Two-kernel approach (correct): + Kernel 1: compute_amax_gsa on the de-interleaved values (GPU-only) + Kernel 2: deinterleave_quantize_from_buffer using gsa from GPU buffer Args: fused_bf16: (M, 2*intermediate) BF16 — fused L1 output @@ -265,8 +266,15 @@ def deinterleave_amax_quantize_nvfp4_fused(fused_bf16, intermediate, divisor=6.0 gsa: (M,) float32 GPU tensor — per-row global scale for L2 GEMM """ from dsv4.kernels.cuda.loader import get_cuda_module - mod = get_cuda_module("fused_deinterleave_amax_quantize", ["fused_deinterleave_amax_quantize.cu"]) - return mod.fused_deinterleave_amax_quantize(fused_bf16, intermediate, granularity, divisor) + # Compute gsa from the fused output + amax_mod = get_cuda_module("amax_gsa", ["amax_gsa.cu"]) + gsa_gpu = amax_mod.compute_amax_gsa(fused_bf16, divisor) + if gsa_gpu.dim() == 0: + gsa_gpu = gsa_gpu.reshape(1) + # Deinterleave + quantize using gsa from GPU buffer + quant_mod = get_cuda_module("fused_amax_quantize", ["fused_amax_quantize.cu"]) + x_fp4, x_sf = quant_mod.deinterleave_quantize_from_buffer(fused_bf16, intermediate, granularity, gsa_gpu) + return x_fp4, x_sf, gsa_gpu def compute_amax_gsa_gpu(x_bf16, divisor=6.0 * 448.0): @@ -284,15 +292,16 @@ def compute_amax_gsa_gpu(x_bf16, divisor=6.0 * 448.0): def quantize_nvfp4_gpu_fused(x_bf16, divisor=6.0 * 448.0): - """Fused amax + gsa + quantize: NO CPU sync, single kernel launch. + """Fused amax + gsa + quantize: zero CPU syncs, two kernel launches. - Replaces the two-step path: - amax = x.float().abs().max().item() ← CPU-GPU sync! - gsa = amax / (6.0 * 448.0) - quantize_nvfp4_gpu(x, gsa) + Two-kernel approach (correct cross-CTA reduction): + Kernel 1: compute_amax_gsa — row-wise amax → gsa on GPU (no .item()) + Kernel 2: quantize_nvfp4_from_buffer — quantize using gsa from GPU buffer - This fused kernel computes amax on GPU, derives gsa, and quantizes - in a single kernel launch. Zero CPU-GPU syncs. + The previous single-kernel approach had a race condition: the cross-CTA + shared memory reduction used __syncthreads() which only syncs within a + CTA, not across CTAs in the same grid. CTA 0 could read s_amax[b] before + CTA b had written it, producing garbage gsa values. Args: x_bf16: (M, N) BF16 tensor. N must be a multiple of 16. @@ -304,8 +313,14 @@ def quantize_nvfp4_gpu_fused(x_bf16, divisor=6.0 * 448.0): gsa: (M,) float32 GPU tensor — per-row global scale for GEMM """ from dsv4.kernels.cuda.loader import get_cuda_module - mod = get_cuda_module("fused_amax_quantize", ["fused_amax_quantize.cu"]) - return mod.fused_amax_quantize_nvfp4(x_bf16, divisor) + amax_mod = get_cuda_module("amax_gsa", ["amax_gsa.cu"]) + gsa_gpu = amax_mod.compute_amax_gsa(x_bf16, divisor) # scalar GPU tensor for M=1 + # Reshape to (M,) for the quantize-from-buffer kernel + if gsa_gpu.dim() == 0: + gsa_gpu = gsa_gpu.reshape(1) # (1,) for M=1 + quant_mod = get_cuda_module("fused_amax_quantize", ["fused_amax_quantize.cu"]) + x_fp4, x_sf = quant_mod.quantize_nvfp4_from_buffer(x_bf16, gsa_gpu) + return x_fp4, x_sf, gsa_gpu def quantize_nvfp4_gpu(x_bf16, global_scale):