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

CRITICAL FIX: fused_amax_quantize cross-CTA race condition

Browse Source 
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.
This commit is contained in:
biondizzle
2026-06-01 21:26:51 +00:00
parent e3412cf913
commit cacf64232e
3 changed files with 167 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
dsv4/kernels/cuda/fused_amax_quantize.cu
View File

@@ -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 <cuda.h>
@@ -40,7 +48,6 @@
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include <cstdint>
#include <cfloat>
__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<torch::Tensor, torch::Tensor, torch::Tensor> fused_amax_quantize_nvfp4_cuda(
torch::Tensor input_bf16, double divisor
// Python API: quantize with gsa from GPU buffer
std::tuple<torch::Tensor, torch::Tensor> 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<<<grid, block, smem_size, c10::cuda::getCurrentCUDAStream()>>>(
quantize_nvfp4_from_buffer_kernel<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<const __nv_bfloat16*>(input_bf16.data_ptr<at::BFloat16>()),
M, N, (float)divisor,
out_fp4.data_ptr<uint8_t>(), out_sf.data_ptr<uint8_t>(),
out_gsa.data_ptr<float>()
M, N, gsa_buffer.data_ptr<float>(),
out_fp4.data_ptr<uint8_t>(), out_sf.data_ptr<uint8_t>()
);
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<torch::Tensor, torch::Tensor> 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<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<const __nv_bfloat16*>(fused_bf16.data_ptr<at::BFloat16>()),
M, N, (int)intermediate, (int)granularity, gsa_buffer.data_ptr<float>(),
out_fp4.data_ptr<uint8_t>(), out_sf.data_ptr<uint8_t>()
);
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);
}

6
dsv4/kernels/cuda/loader.py
View File

@@ -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"])

43
dsv4/ops/quantize.py
View File

@@ -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):