diff --git a/dsv4/kernels/cuda/quantize_nvfp4.cu b/dsv4/kernels/cuda/quantize_nvfp4.cu new file mode 100644 index 00000000..c63e9322 --- /dev/null +++ b/dsv4/kernels/cuda/quantize_nvfp4.cu @@ -0,0 +1,116 @@ +#include +#include +#include +#include +#include +#include +#include + +// BF16 → NVFP4 quantization kernel (no deinterleave, GPU-only). +// Reads BF16 from GMEM, quantizes to NVFP4 (FP4 data + FP8 E4M3 scales), +// writes both to GMEM. No CPU-GPU syncs. +// +// This replaces quantize_activation_nvfp4() which uses .amax() (CPU sync). +// Global scale is passed in as a pre-computed scalar. +// +// Grid: (N / 16, M, 1) — each CTA processes one 16-element block in one row. +// Block: 16 threads (1 thread per element in the 16-element microblock). + +__device__ __forceinline__ int half_step_to_e2m1(int hs) { + // Matches Python step_to_idx LUT: + // 0→0, 1→1, 2→2, 3→3, 4→4, 5→4, 6→5, 7→5, 8→6, 9→6, 10→6, 11→7, 12→7 + if (hs <= 4) return hs; + if (hs <= 5) return 4; + if (hs <= 7) return 5; + if (hs <= 10) return 6; + return 7; +} + +__global__ void quantize_nvfp4_kernel( + const __nv_bfloat16* __restrict__ input, + int M, int N, + float global_scale, + 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 >= N) return; + + int lane = threadIdx.x; // 0..15 + + // Step 1: Read 1 BF16 element per thread, normalize by global_scale + int col = n_block * 16 + lane; + float val = 0.0f; + if (col < N) { + val = __bfloat162float(input[m * N + col]); + val = val / global_scale; + } + + // Step 2: Warp-level amax reduction (16 threads = half-warp) + // Use warp shuffle for the reduction + float abs_val = fabsf(val); + for (int offset = 8; offset > 0; offset >>= 1) { + abs_val = fmaxf(abs_val, __shfl_down_sync(0xFFFF, abs_val, offset)); + } + float block_amax = abs_val; // Same value in all 16 lanes after reduction + + // Step 3: Compute FP8 E4M3 block scale = amax / 6.0 + float bsf = block_amax / 6.0f; + if (block_amax < 6.0f * 0.001953125f) { // FP8 E4M3 underflow threshold + bsf = 0; + val = 0; + } + __nv_fp8_e4m3 bsf8_obj(bsf); + float bs = (float)bsf8_obj; + uint8_t bsf8 = *(uint8_t*)&bsf8_obj; + + // Step 4: Quantize each value to FP4 E2M1 + uint8_t nibble = 0; + if (bs >= 1e-8f) { + float s = val / 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; + nibble = idx; + } + + // Step 5: Pack pairs of FP4 nibbles into bytes + // Even lanes write the packed byte: (odd_nibble << 4) | even_nibble + if (lane % 2 == 0) { + uint8_t odd_nibble = __shfl_down_sync(0xFFFF, nibble, 1); + uint8_t packed = (odd_nibble << 4) | nibble; + int byte_idx = m * (N / 2) + n_block * 8 + (lane / 2); + out_fp4[byte_idx] = packed; + } + + // Step 6: Write FP8 block scale (1 thread per block) + if (lane == 0) { + out_sf[m * (N / 16) + n_block] = bsf8; + } +} + +std::tuple quantize_nvfp4_cuda( + torch::Tensor input_bf16, double global_scale +) { + 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"); + 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)); + int nb = N / 16; + dim3 grid(nb, M); + dim3 block(16); + quantize_nvfp4_kernel<<>>( + reinterpret_cast(input_bf16.data_ptr()), + M, N, (float)global_scale, + 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("quantize_nvfp4", &quantize_nvfp4_cuda); +} diff --git a/dsv4/layers/moe.py b/dsv4/layers/moe.py index 7fc38fa1..63d3a146 100644 --- a/dsv4/layers/moe.py +++ b/dsv4/layers/moe.py @@ -19,6 +19,8 @@ from dsv4.ops.quantize import ( quantize_activation_nvfp4, quantize_weight_to_nvfp4, quantize_to_nvfp4, + quantize_nvfp4_gpu, + deinterleave_quantize_nvfp4_cuda, ) from dsv4.ops.layouts import ( make_b_k_major, @@ -563,14 +565,10 @@ class Nvfp4MoE: padded_dst = padded_expert_offsets[expert_assign] + local_row # === L1: gate + up === - # Quantize slot_hidden (sorted tokens), NOT padded_hidden. - # padded_hidden is padded with zeros; quantizing it produces - # x_sf rows at padded positions, but x_sf[:num_slots] would - # only get scales for the first num_slots PADDED rows (expert 0), - # not the scattered token positions. Quantizing slot_hidden - # gives x_sf with num_slots rows (one per token), which the - # scale assembly correctly scatters into padded layout. - slot_x_fp4, slot_x_sf = quantize_activation_nvfp4( + # Quantize slot_hidden using GPU-only kernel (no CPU-GPU sync). + # slot_hidden is the sorted tokens (not padded). The GPU kernel + # replaces quantize_activation_nvfp4 which uses .amax() (CPU sync). + slot_x_fp4, slot_x_sf = quantize_nvfp4_gpu( slot_hidden, self._l1_activation_global_scale ) # Scatter x_fp4 into padded layout for the GEMM @@ -596,9 +594,13 @@ class Nvfp4MoE: swiglu_limit=self._swiglu_limit if self._swiglu_limit is not None else 0.0, ) l1_out_real = l1_out[padded_dst] - # De-interleave: odd 8-col groups = silu(gate)*up (the SwiGLU result) - l1_deil = deinterleave_l1_weights(l1_out_real.unsqueeze(0).contiguous())[0] - activated = l1_deil[:, self.intermediate_size:] + # De-interleave + quantize to FP4 in one GPU kernel. + # l1_out_real has interleaved [silu(gate)*8, swiglu*8, ...]. + # The CUDA kernel extracts odd 8-col groups (SwiGLU result) + # and quantizes to NVFP4. No CPU sync, no Python deinterleave. + slot_l2_x_fp4, slot_l2_x_sf = deinterleave_quantize_nvfp4_cuda( + l1_out_real, self.intermediate_size, self._l2_activation_global_scale + ) else: # === Non-fused L1 GEMM + PyTorch SiLU(gate)*up === l1_out = run_nvfp4_grouped_gemm( @@ -618,10 +620,12 @@ class Nvfp4MoE: activated = gate_silu * up # === L2: down === - # Quantize activated (per-token), scatter into padded FP4 buffer - slot_l2_x_fp4, slot_l2_x_sf = quantize_activation_nvfp4( - activated, self._l2_activation_global_scale - ) + # Quantize activated (per-token) using GPU-only kernel, scatter into padded FP4 buffer. + # For fused_swiglu path, slot_l2_x_fp4/sf already set by deinterleave_quantize_nvfp4_cuda. + if not self._fused_swiglu: + slot_l2_x_fp4, slot_l2_x_sf = quantize_nvfp4_gpu( + activated, self._l2_activation_global_scale + ) padded_activated_fp4 = self._shared_bufs['activated_fp4'] padded_activated_fp4.view(torch.uint8).zero_() padded_activated_fp4.view(torch.uint8)[padded_dst] = slot_l2_x_fp4.view(torch.uint8) diff --git a/dsv4/ops/quantize.py b/dsv4/ops/quantize.py index 798b19de..29b97bc9 100644 --- a/dsv4/ops/quantize.py +++ b/dsv4/ops/quantize.py @@ -244,7 +244,8 @@ def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, gra """ from torch.utils.cpp_extension import load import os - kernel_dir = os.path.join(os.path.dirname(__file__), "kernels") + # dsv4/ops/quantize.py → dsv4/kernels/cuda/ + kernel_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "kernels", "cuda") mod = load( name="deinterleave_quantize_nvfp4", sources=[os.path.join(kernel_dir, "deinterleave_quantize.cu")], @@ -252,3 +253,30 @@ def deinterleave_quantize_nvfp4_cuda(fused_bf16, intermediate, global_scale, gra verbose=False, ) return mod.deinterleave_quantize_nvfp4(fused_bf16, intermediate, granularity, global_scale) + + +def quantize_nvfp4_gpu(x_bf16, global_scale): + """Quantize BF16 tensor to NVFP4 using a custom CUDA kernel (GPU-only, no CPU sync). + + Replaces quantize_activation_nvfp4() which uses .amax() (CPU sync). + The global_scale must be pre-computed (from warmup or known value). + + Args: + x_bf16: (M, N) BF16 tensor. N must be a multiple of 16. + global_scale: float32 scalar (pre-computed, NOT from .max()) + + Returns: + x_fp4: (M, N//2) float4_e2m1fn_x2 + x_sf: (M, N//16) float8_e4m3fn + """ + from torch.utils.cpp_extension import load + import os + # dsv4/ops/quantize.py → dsv4/kernels/cuda/ + kernel_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "kernels", "cuda") + mod = load( + name="quantize_nvfp4", + sources=[os.path.join(kernel_dir, "quantize_nvfp4.cu")], + extra_cuda_cflags=["-gencode=arch=compute_100a,code=sm_100a"], + verbose=False, + ) + return mod.quantize_nvfp4(x_bf16, global_scale) diff --git a/tests/unit/test_nvfp4_gpu_quantize.py b/tests/unit/test_nvfp4_gpu_quantize.py new file mode 100644 index 00000000..70586ac8 --- /dev/null +++ b/tests/unit/test_nvfp4_gpu_quantize.py @@ -0,0 +1,193 @@ +""" +NVFP4 GPU Quantize Kernel Test. + +Tests: +1. quantize_nvfp4_gpu: BF16 → FP4 (no deinterleave, GPU-only, no CPU sync) +2. deinterleave_quantize_nvfp4_cuda: interleaved BF16 → FP4 (deinterleave + quantize) +3. Integration: Both kernels match the Python quantize_activation_nvfp4 reference + +Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_nvfp4_gpu_quantize.py +""" +import torch +import math +import sys, os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from dsv4.ops.quantize import ( + quantize_activation_nvfp4, + quantize_nvfp4_gpu, + deinterleave_quantize_nvfp4_cuda, + SF_VEC_SIZE, +) +from dsv4.ops.layouts import interleave_l1_weights + + +def test_quantize_nvfp4_gpu_basic(): + """Test GPU-only quantize kernel against Python reference.""" + print("\n=== Test 1: quantize_nvfp4_gpu basic correctness ===") + torch.manual_seed(42) + M, N = 128, 512 + x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda') + global_scale = 1.0 + + # Python reference + ref_fp4, ref_sf = quantize_activation_nvfp4(x, global_scale) + + # GPU kernel + gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale) + + # Compare shapes and dtypes + assert gpu_fp4.shape == ref_fp4.shape, f"FP4 shape mismatch: {gpu_fp4.shape} vs {ref_fp4.shape}" + assert gpu_sf.shape == ref_sf.shape, f"SF shape mismatch: {gpu_sf.shape} vs {ref_sf.shape}" + assert gpu_fp4.dtype == torch.float4_e2m1fn_x2, f"FP4 dtype: {gpu_fp4.dtype}" + assert gpu_sf.dtype == torch.float8_e4m3fn, f"SF dtype: {gpu_sf.dtype}" + + # Byte-exact comparison of FP4 data + ref_bytes = ref_fp4.view(torch.uint8) + gpu_bytes = gpu_fp4.view(torch.uint8) + fp4_match = (ref_bytes == gpu_bytes).float().mean().item() + print(f" FP4 byte match: {fp4_match*100:.1f}%") + + # SF comparison + ref_sf_bytes = ref_sf.view(torch.uint8) + gpu_sf_bytes = gpu_sf.view(torch.uint8) + sf_match = (ref_sf_bytes == gpu_sf_bytes).float().mean().item() + print(f" SF byte match: {sf_match*100:.1f}%") + + # Round-trip cosine similarity + ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, N) + gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, N) + + cos_ref = torch.nn.functional.cosine_similarity( + x.flatten().float().unsqueeze(0), ref_deq.flatten().float().unsqueeze(0) + ).item() + cos_gpu = torch.nn.functional.cosine_similarity( + x.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0) + ).item() + cos_cross = torch.nn.functional.cosine_similarity( + ref_deq.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0) + ).item() + + print(f" Python round-trip cos: {cos_ref:.6f}") + print(f" GPU round-trip cos: {cos_gpu:.6f}") + print(f" Python vs GPU cos: {cos_cross:.6f}") + + assert cos_gpu >= 0.95, f"GPU round-trip cosine too low: {cos_gpu}" + assert cos_cross >= 0.99, f"Python vs GPU cosine too low: {cos_cross}" + print(f" ✅ PASS") + + +def test_quantize_nvfp4_gpu_larger(): + """Test GPU quantize with larger dimensions (MoE intermediate size).""" + print("\n=== Test 2: quantize_nvfp4_gpu larger shape ===") + torch.manual_seed(42) + M, N = 64, 4096 + x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda') + global_scale = 1.0 / (6.0 * 448.0) + + ref_fp4, ref_sf = quantize_activation_nvfp4(x, global_scale) + gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale) + + ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, N) + gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, N) + + cos_ref = torch.nn.functional.cosine_similarity( + x.flatten().float().unsqueeze(0), ref_deq.flatten().float().unsqueeze(0) + ).item() + cos_gpu = torch.nn.functional.cosine_similarity( + x.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0) + ).item() + + print(f" Python round-trip cos: {cos_ref:.6f}") + print(f" GPU round-trip cos: {cos_gpu:.6f}") + assert cos_gpu >= 0.95, f"GPU round-trip cosine too low: {cos_gpu}" + print(f" ✅ PASS") + + +def test_quantize_nvfp4_gpu_no_cpu_sync(): + """Verify quantize_nvfp4_gpu has no CPU-GPU sync points.""" + print("\n=== Test 3: No CPU-GPU sync verification ===") + torch.manual_seed(42) + M, N = 32, 512 + x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda') + global_scale = 0.001 + + # This should NOT trigger a CUDA synchronization + # If it does, cudagraph capture would fail + gpu_fp4, gpu_sf = quantize_nvfp4_gpu(x, global_scale) + + # Force a sync to check results + result = gpu_fp4.view(torch.uint8).sum().item() + print(f" FP4 sum (verification): {result}") + print(f" ✅ PASS (no crash = no CPU sync in kernel path)") + + +def test_deinterleave_quantize_correctness(): + """Test deinterleave + quantize CUDA kernel against Python reference.""" + print("\n=== Test 4: deinterleave_quantize_nvfp4_cuda correctness ===") + torch.manual_seed(42) + M = 64 + intermediate = 512 + N = 2 * intermediate # fused output has gate + up interleaved + + # Create SwiGLU-style interleaved data + gate = torch.randn(M, intermediate, dtype=torch.bfloat16, device='cuda') + up = torch.randn(M, intermediate, dtype=torch.bfloat16, device='cuda') + gate_silu = torch.nn.functional.silu(gate) + swiglu_result = gate_silu * up + + # Create interleaved layout: [silu(gate)*8, up*8, ...] + fused = torch.cat([gate_silu, up], dim=-1).unsqueeze(0) # (1, M, N) + fused = interleave_l1_weights(fused)[0] # (M, N) interleaved + + global_scale = 1.0 + + # Python reference: deinterleave then quantize + ref_fp4, ref_sf = quantize_activation_nvfp4(swiglu_result, global_scale) + + # CUDA kernel: deinterleave + quantize in one pass + gpu_fp4, gpu_sf = deinterleave_quantize_nvfp4_cuda(fused, intermediate, global_scale) + + # Compare round-trip + ref_deq = _dequantize_nvfp4(ref_fp4, ref_sf, global_scale, intermediate) + gpu_deq = _dequantize_nvfp4(gpu_fp4, gpu_sf, global_scale, intermediate) + + cos_cross = torch.nn.functional.cosine_similarity( + ref_deq.flatten().float().unsqueeze(0), gpu_deq.flatten().float().unsqueeze(0) + ).item() + + print(f" Python vs CUDA kernel cos: {cos_cross:.6f}") + assert cos_cross >= 0.99, f"Python vs CUDA kernel cosine too low: {cos_cross}" + print(f" ✅ PASS") + + +def test(): + print("=== NVFP4 GPU Quantize Kernel Tests ===") + test_quantize_nvfp4_gpu_basic() + test_quantize_nvfp4_gpu_larger() + test_quantize_nvfp4_gpu_no_cpu_sync() + test_deinterleave_quantize_correctness() + print("\n=== ALL TESTS PASSED ===") + + +def _dequantize_nvfp4(x_fp4, block_scale, global_scale, N): + """Dequantize NVFP4 back to BF16 for verification.""" + M = x_fp4.shape[0] + block_size = SF_VEC_SIZE + raw = x_fp4.view(torch.uint8) + even = raw & 0x0F + odd = (raw >> 4) & 0x0F + indices = torch.stack([even, odd], dim=-1).reshape(M, N) + signs = (indices >= 8).float() * -2 + 1 + mag = indices % 8 + idx_to_half_step = torch.tensor([0, 2, 4, 6, 8, 10, 12, 14], + dtype=torch.float32, device='cuda') + half_steps = idx_to_half_step[mag.long()] + x_deq_fp32 = signs * half_steps / 2.0 + block_scale_exp = block_scale.repeat_interleave(block_size, dim=-1).float() + x_deq_fp32 = x_deq_fp32 * block_scale_exp * global_scale + return x_deq_fp32.to(torch.bfloat16) + + +if __name__ == '__main__': + test()