185 lines
6.3 KiB
Python
185 lines
6.3 KiB
Python
"""
|
|
NVFP4-1.1 Phase 1: Verify FP4 quantization math in CuTeDSL kernel.
|
|
|
|
Tests that fp4_quant.py functions produce bit-exact matches with the
|
|
Python reference (quantize_activation_nvfp4). Runs on B200 only.
|
|
|
|
Strategy: Launch a kernel that processes 16 BF16 values through the
|
|
quantization pipeline and writes results to GMEM. Compare with Python.
|
|
|
|
Uses cute.arch.load for scalar GMEM reads (proven pattern).
|
|
For GMEM writes, uses cute.copy with a simple CopyUniversalOp atom.
|
|
"""
|
|
|
|
import torch
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.torch as cutlass_torch
|
|
from cutlass.cute.nvgpu import cpasync
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))
|
|
|
|
from dsv4.ops.quantize import quantize_activation_nvfp4, SF_VEC_SIZE
|
|
from dsv4.kernels.gemm.fp4_quant import (
|
|
fp8_e4m3_from_float32,
|
|
fp8_e4m3_to_float32,
|
|
half_step_to_e2m1_idx,
|
|
quantize_e2m1_nibble,
|
|
)
|
|
|
|
|
|
@cute.kernel
|
|
def fp4_quant_test_kernel(
|
|
input_bf16: cute.Tensor, # (16,) BF16 — 16 input values
|
|
out_fp4: cute.Tensor, # (8,) Int32 — packed FP4 bytes
|
|
out_sf: cute.Tensor, # (1,) Int32 — FP8 scale byte
|
|
gs_scalar: cute.Tensor, # (1,) Float32 — global scale
|
|
):
|
|
"""Quantize 16 BF16 values to NVFP4 using fp4_quant functions.
|
|
|
|
Single-thread kernel (only thread 0 does work).
|
|
Grid: (1, 1, 1), Block: (32, 1, 1)
|
|
"""
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
|
|
if tidx == cutlass.Int32(0):
|
|
# Load global scale
|
|
gs = cute.arch.load(gs_scalar.iterator, cutlass.Float32)
|
|
|
|
# Load 16 BF16 values, convert to FP32, normalize by global_scale
|
|
vals_f32 = [cutlass.Float32(0.0)] * 16
|
|
for i in cutlass.range(16, unroll=1):
|
|
bf16_val = cute.arch.load(
|
|
input_bf16.iterator + i * cutlass.Int32(2), # BF16 = 2 bytes
|
|
cutlass.BFloat16,
|
|
)
|
|
vals_f32[i] = bf16_val.to(cutlass.Float32) / gs
|
|
|
|
# Compute per-16-element amax
|
|
amax = cutlass.Float32(0.0)
|
|
for i in cutlass.range(16, unroll=1):
|
|
v = vals_f32[i]
|
|
a = cute.math.fmax(v, cutlass.Float32(0.0) - v) # abs
|
|
amax = cute.math.fmax(amax, a)
|
|
|
|
# Block scale = amax / 6
|
|
bsf_f32 = amax / cutlass.Float32(6.0)
|
|
# Underflow: if amax < 6 * 2^-9, force scale = 0
|
|
underflow_threshold = cutlass.Float32(6.0 * (2.0 ** -9))
|
|
if amax < underflow_threshold:
|
|
bsf_f32 = cutlass.Float32(0.0)
|
|
|
|
# FP8 E4M3 cast
|
|
sf_bits = fp8_e4m3_from_float32(bsf_f32)
|
|
|
|
# Dequantize FP8 scale (round-trip)
|
|
bs_dequant = fp8_e4m3_to_float32(sf_bits)
|
|
|
|
# Quantize each value to E2M1 and pack
|
|
for i in cutlass.range(8, unroll=1):
|
|
nibble0 = quantize_e2m1_nibble(vals_f32[2 * i], bs_dequant)
|
|
nibble1 = quantize_e2m1_nibble(vals_f32[2 * i + 1], bs_dequant)
|
|
packed = (nibble1 << cutlass.Int32(4)) | nibble0
|
|
# Write packed byte as Int32
|
|
cute.arch.store(out_fp4.iterator + i * cutlass.Int32(4), packed, cutlass.Int32)
|
|
|
|
# Write FP8 scale byte
|
|
cute.arch.store(out_sf.iterator, sf_bits, cutlass.Int32)
|
|
|
|
|
|
def run_test():
|
|
"""Run the FP4 quantization test on GPU."""
|
|
device = "cuda"
|
|
N = 16
|
|
|
|
# Generate test input
|
|
torch.manual_seed(42)
|
|
x_bf16 = torch.randn(1, N, dtype=torch.bfloat16, device=device)
|
|
|
|
# Compute global scale
|
|
x_f32 = x_bf16.float()
|
|
amax_val = x_f32.abs().max().item()
|
|
global_scale = max(amax_val / (6.0 * 448.0), 1e-8)
|
|
|
|
# Python reference
|
|
ref_fp4, ref_sf = quantize_activation_nvfp4(x_bf16, global_scale)
|
|
ref_fp4_bytes = ref_fp4.view(torch.uint8).reshape(-1).cpu()
|
|
ref_sf_bytes = ref_sf.view(torch.uint8).cpu()
|
|
|
|
print(f"Input BF16 (first 8): {x_bf16[0, :8].cpu()}")
|
|
print(f"Global scale: {global_scale:.8f}")
|
|
print(f"Ref FP4: {ref_fp4_bytes}")
|
|
print(f"Ref SF: {ref_sf_bytes}")
|
|
|
|
# Prepare output tensors
|
|
out_fp4 = torch.zeros(8, dtype=torch.int32, device=device)
|
|
out_sf = torch.zeros(1, dtype=torch.int32, device=device)
|
|
gs_tensor = torch.tensor([global_scale], dtype=torch.float32, device=device)
|
|
|
|
# Convert to CuTe tensors
|
|
def to_cute(t):
|
|
ct = cutlass_torch.from_dlpack(t)
|
|
return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t))
|
|
|
|
x_flat = x_bf16.reshape(N).contiguous()
|
|
input_c = to_cute(x_flat)
|
|
out_fp4_c = to_cute(out_fp4)
|
|
out_sf_c = to_cute(out_sf)
|
|
gs_c = to_cute(gs_tensor)
|
|
|
|
# Compile and run
|
|
import cuda.bindings.driver as cuda
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
|
|
print("\nCompiling kernel (first run may take a minute)...")
|
|
try:
|
|
compiled = cute.compile(
|
|
fp4_quant_test_kernel,
|
|
input_c, out_fp4_c, out_sf_c, gs_c,
|
|
stream,
|
|
)
|
|
print("Compiled. Running...")
|
|
compiled(input_c, out_fp4_c, out_sf_c, gs_c, stream)
|
|
torch.cuda.synchronize()
|
|
|
|
# Extract results
|
|
our_fp4 = out_fp4[:8].to(torch.uint8).cpu()
|
|
our_sf = out_sf[0].to(torch.uint8).cpu().item()
|
|
|
|
print(f"\nOur FP4: {our_fp4}")
|
|
print(f"Our SF: {our_sf}")
|
|
|
|
fp4_match = torch.equal(our_fp4, ref_fp4_bytes[:8])
|
|
sf_match = our_sf == ref_sf_bytes[0].item()
|
|
|
|
if fp4_match and sf_match:
|
|
print("\n✅ PASS: FP4 quantization matches Python reference!")
|
|
return True
|
|
else:
|
|
print(f"\n❌ FAIL: FP4 match={fp4_match}, SF match={sf_match}")
|
|
if not fp4_match:
|
|
for i in range(8):
|
|
o = our_fp4[i].item()
|
|
r = ref_fp4_bytes[i].item()
|
|
if o != r:
|
|
print(f" Byte {i}: ours=0x{o:02x}, ref=0x{r:02x}")
|
|
if not sf_match:
|
|
print(f" SF: ours=0x{our_sf:02x}, ref=0x{ref_sf_bytes[0].item():02x}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"\n❌ ERROR: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 60)
|
|
print("NVFP4-1.1 Phase 1: FP4 Quantization Math Test")
|
|
print("Verifies fp4_quant.py functions match Python reference")
|
|
print("=" * 60)
|
|
success = run_test()
|
|
exit(0 if success else 1)
|