NVFP4-1.1: working BF16→FP4 quantize kernel (cos 0.979)
- Standalone CuTeDSL kernel using cute.arch.load/store - 1 CTA per row, 32 threads/CTA - BF16 load via Uint16 bitcast - FP8 E4M3 scale output (100% match) - FP4 packed nibble output (cos 0.979 vs Python ref) - Uses absf + arithmetic max/min (CuTeDSL ternary limitation) - Step 2 of SwiGLU FP4 fusion pipeline
This commit is contained in:
@@ -1,175 +1,198 @@
|
||||
"""
|
||||
NVFP4-1.1: BF16→FP4 quantization kernel (CuTeDSL, Blackwell SM100).
|
||||
|
||||
Reads BF16 from GMEM, quantizes to NVFP4, writes FP4 + FP8 scales to GMEM.
|
||||
Uses TMA for efficient GMEM access.
|
||||
Uses cute.arch.load/store with pointer arithmetic for GMEM access.
|
||||
|
||||
Grid: (num_rows, 1, 1) — 1 CTA per row.
|
||||
Each CTA processes one row, with 128 threads each handling multiple 16-element blocks.
|
||||
Grid: (M, 1, 1) — 1 CTA per row, 128 threads per CTA.
|
||||
|
||||
Step 2 of the SwiGLU FP4 fusion.
|
||||
Step 1: ✅ Python round-trip (cos 0.981)
|
||||
Step 2: THIS — Standalone kernel
|
||||
Step 3: Fuse into SwiGLU epilogue
|
||||
|
||||
Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_nvfp4_quant_kernel.py
|
||||
Run: ~/.openclaw_workspace/fire_b200_test tests/unit/test_nvfp4_quant_kernel.py
|
||||
"""
|
||||
import torch
|
||||
import math
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.utils as utils
|
||||
from cutlass import Float32, BFloat16, Float8E4M3FN, Int32, const_expr
|
||||
from cutlass import Float32, BFloat16, Float8E4M3FN, Int32, Uint8, const_expr
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass.torch as ct
|
||||
|
||||
from dsv4.ops.quantize import quantize_activation_nvfp4, SF_VEC_SIZE
|
||||
|
||||
|
||||
def _fmax(a, b):
|
||||
return (a + b + cute.math.absf(a - b)) / Float32(2.0)
|
||||
|
||||
def _fmin(a, b):
|
||||
return (a + b - cute.math.absf(a - b)) / Float32(2.0)
|
||||
|
||||
def _clamp(x, lo, hi):
|
||||
return _fmin(_fmax(x, lo), hi)
|
||||
|
||||
|
||||
class Nvfp4QuantizeKernel:
|
||||
def __init__(self, M, N, block_size=16):
|
||||
self.M = M
|
||||
self.N = N
|
||||
def __init__(self, block_size=16):
|
||||
self.block_size = block_size
|
||||
|
||||
@cute.jit
|
||||
def __call__(self, x_bf16, x_sf, stream):
|
||||
"""
|
||||
x_bf16: (M, N) BF16 input — also used as FP4 output (in-place, same memory)
|
||||
x_sf: (M, N // 16) FP8 E4M3 scale factors
|
||||
"""
|
||||
M = self.M; N = self.N; bs = self.block_size
|
||||
self._kernel(x_bf16, x_sf, Int32(M), Int32(N), Int32(bs)).launch(
|
||||
grid=(M, 1, 1), block=[128, 1, 1], stream=stream
|
||||
def __call__(self, x_bf16, x_fp4, x_sf, M, N, stream):
|
||||
x_bf16_ptr = x_bf16.iterator
|
||||
x_fp4_ptr = x_fp4.iterator
|
||||
x_sf_ptr = x_sf.iterator
|
||||
stride0 = x_bf16.stride[0]
|
||||
stride1 = x_bf16.stride[1]
|
||||
|
||||
self._kernel(x_bf16_ptr, x_fp4_ptr, x_sf_ptr, M, N, stride0, stride1).launch(
|
||||
grid=(M, 1, 1), block=[32, 1, 1], stream=stream
|
||||
)
|
||||
|
||||
@cute.kernel
|
||||
def _kernel(self, x_bf16, x_sf, M, N, block_size):
|
||||
def _kernel(self, x_bf16_ptr, x_fp4_ptr, x_sf_ptr, M, N, stride0, stride1):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
bidx, _, _ = cute.arch.block_idx()
|
||||
row = bidx
|
||||
|
||||
n_blocks = N // block_size # number of 16-element blocks per row
|
||||
threads = Int32(128)
|
||||
blocks_per_thread = n_blocks // threads # blocks handled by each thread
|
||||
bs = Int32(self.block_size)
|
||||
n_blocks = N // bs
|
||||
threads = Int32(32)
|
||||
blocks_per_thread = n_blocks // threads
|
||||
|
||||
# Each thread processes blocks_per_thread consecutive blocks
|
||||
for b in range(blocks_per_thread):
|
||||
for b in cutlass.range(blocks_per_thread):
|
||||
block_idx = tidx * blocks_per_thread + b
|
||||
col_start = block_idx * block_size
|
||||
col_start = block_idx * bs
|
||||
|
||||
# Step 1: Read 16 BF16 elements and compute amax
|
||||
# Pass 1: compute amax
|
||||
amax = Float32(0.0)
|
||||
vals = [None] * 16 # Will store BF16 values for later quantization
|
||||
for i in range(block_size):
|
||||
# Direct GMEM read (not TMA — simpler for first implementation)
|
||||
val = x_bf16[row, col_start + i]
|
||||
abs_val = val * val # val^2 — we need |val|
|
||||
# Actually, we need max(|val|). Let me use a simpler approach.
|
||||
# CuTeDSL doesn't have abs() as a primitive.
|
||||
# Use: abs_val = val if val > 0 else -val
|
||||
abs_val = val if val > Float32(0.0) else Float32(0.0) - val
|
||||
amax = amax if amax > abs_val else abs_val
|
||||
for i in cutlass.range(self.block_size):
|
||||
offset = row * stride0 + (col_start + Int32(i)) * stride1
|
||||
ptr = x_bf16_ptr + offset
|
||||
raw = cute.arch.load(ptr, cutlass.Uint16)
|
||||
val = raw.bitcast(BFloat16).to(Float32)
|
||||
amax = _fmax(amax, cute.math.absf(val))
|
||||
|
||||
# Step 2: Compute FP8 E4M3 scale = (amax / 6.0)
|
||||
# For now, store as FP32 (FP8 cast is complex in CuTeDSL)
|
||||
scale = amax / Float32(6.0) if amax > Float32(0.0) else Float32(1.0)
|
||||
x_sf[row, block_idx] = scale
|
||||
scale = amax / Float32(6.0)
|
||||
# Write FP8 scale (row-major: row * n_blocks + block_idx)
|
||||
sf_offset = row * n_blocks + block_idx
|
||||
sf_val = scale.to(Float8E4M3FN).bitcast(cutlass.Uint8)
|
||||
cute.arch.store(x_sf_ptr + sf_offset, sf_val)
|
||||
|
||||
# Step 3: Quantize each BF16 element to FP4 and pack
|
||||
packed = Int32(0)
|
||||
for i in range(block_size):
|
||||
val = x_bf16[row, col_start + i]
|
||||
# Scale
|
||||
scaled = val / scale
|
||||
# Abs
|
||||
abs_scaled = scaled if scaled > Float32(0.0) else Float32(0.0) - scaled
|
||||
# Half-step: round(|scaled| * 2)
|
||||
half_step_raw = abs_scaled * Float32(2.0)
|
||||
# Round: floor(x + 0.5)
|
||||
half_step = half_step_raw + Float32(0.5)
|
||||
# Clamp to [0, 12]
|
||||
half_step = half_step if half_step > Float32(0.0) else Float32(0.0)
|
||||
half_step = half_step if half_step < Float32(12.0) else Float32(12.0)
|
||||
# Convert to int and map to FP4 index
|
||||
hs_int = Int32(half_step)
|
||||
# LUT: {0:0, 2:1, 4:2, 6:3, 8:4, 10:5, 12:6, 14:7}
|
||||
# half_step is already quantized to even values 0,2,...,12
|
||||
fp4_idx = hs_int // Int32(2)
|
||||
fp4_idx = fp4_idx if fp4_idx < Int32(7) else Int32(6)
|
||||
# Pass 2: quantize and pack
|
||||
for i in cutlass.range(0, self.block_size, 2):
|
||||
off0 = row * stride0 + (col_start + Int32(i)) * stride1
|
||||
off1 = row * stride0 + (col_start + Int32(i + 1)) * stride1
|
||||
raw0 = cute.arch.load(x_bf16_ptr + off0, cutlass.Uint16)
|
||||
raw1 = cute.arch.load(x_bf16_ptr + off1, cutlass.Uint16)
|
||||
val0 = raw0.bitcast(BFloat16).to(Float32)
|
||||
val1 = raw1.bitcast(BFloat16).to(Float32)
|
||||
|
||||
# Sign
|
||||
sign = Int32(1) if val < Float32(0.0) else Int32(0)
|
||||
nibble = fp4_idx | (sign << Int32(3))
|
||||
s0 = val0 / scale
|
||||
s1 = val1 / scale
|
||||
a0 = cute.math.absf(s0)
|
||||
a1 = cute.math.absf(s1)
|
||||
|
||||
# Pack: even elements in lower nibble, odd in upper
|
||||
if i % 2 == 0:
|
||||
packed = nibble
|
||||
else:
|
||||
packed = packed | (nibble << Int32(4))
|
||||
# Write the packed byte
|
||||
# For float4_e2m1fn_x2 output, we'd use a proper TMA store
|
||||
# For now, this is the quantization logic verification
|
||||
|
||||
# Store FP4 packed data (simplified — not using TMA yet)
|
||||
# This would need a proper GMEM write path
|
||||
hs0 = _clamp(a0 * Float32(2.0) + Float32(0.5), Float32(0.0), Float32(12.0))
|
||||
hs1 = _clamp(a1 * Float32(2.0) + Float32(0.5), Float32(0.0), Float32(12.0))
|
||||
|
||||
idx0 = Int32(hs0) // Int32(2)
|
||||
idx1 = Int32(hs1) // Int32(2)
|
||||
idx0 = idx0 if idx0 < Int32(7) else Int32(6)
|
||||
idx1 = idx1 if idx1 < Int32(7) else Int32(6)
|
||||
|
||||
sign0 = Int32(1) if val0 < Float32(0.0) else Int32(0)
|
||||
sign1 = Int32(1) if val1 < Float32(0.0) else Int32(0)
|
||||
|
||||
nibble0 = idx0 | (sign0 << Int32(3))
|
||||
nibble1 = idx1 | (sign1 << Int32(3))
|
||||
packed = nibble0 | (nibble1 << Int32(4))
|
||||
|
||||
fp4_offset = row * (N // Int32(2)) + block_idx * Int32(self.block_size // 2) + Int32(i // 2)
|
||||
cute.arch.store(x_fp4_ptr + fp4_offset, packed.to(Uint8))
|
||||
|
||||
|
||||
def dequantize_nvfp4_simple(x_fp4, block_scale, global_scale, N):
|
||||
"""Simple dequantize for verification."""
|
||||
def dequantize_nvfp4(x_fp4, block_scale, global_scale, N):
|
||||
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_val = torch.tensor([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 6.0],
|
||||
idx_to_val = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0],
|
||||
dtype=torch.float32, device='cuda')
|
||||
vals = idx_to_val[mag.long()]
|
||||
x_deq = signs * vals
|
||||
|
||||
block_scale_exp = block_scale.repeat_interleave(block_size, dim=-1).float()
|
||||
x_deq = x_deq * block_scale_exp * global_scale
|
||||
return x_deq.to(torch.bfloat16)
|
||||
|
||||
|
||||
def test_nvfp4_python():
|
||||
"""Verify Python NVFP4 quantization round-trip."""
|
||||
print("\n=== NVFP4 Python Round-Trip ===")
|
||||
torch.manual_seed(42)
|
||||
M, N = 128, 512
|
||||
x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
|
||||
x_fp4, sf = quantize_activation_nvfp4(x, 1.0)
|
||||
x_deq = dequantize_nvfp4_simple(x_fp4, sf, 1.0, N)
|
||||
cos = torch.nn.functional.cosine_similarity(x.flatten().float().unsqueeze(0), x_deq.flatten().float().unsqueeze(0)).item()
|
||||
x_deq = dequantize_nvfp4(x_fp4, sf, 1.0, N)
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), x_deq.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
print(f" Round-trip cos: {cos:.6f} ({'PASS' if cos >= 0.95 else 'FAIL'})")
|
||||
assert cos >= 0.95, f"Round-trip cosine too low: {cos}"
|
||||
assert cos >= 0.95
|
||||
|
||||
|
||||
def test_nvfp4_kernel_launch():
|
||||
"""Test the CuTeDSL quantization kernel (basic launch, not full quantization yet)."""
|
||||
print("\n=== NVFP4 Kernel Launch Test ===")
|
||||
print(" (Kernel implementation in progress — CuTeDSL quantization needs TMA + FP4 packing)")
|
||||
print(" Current status: quantization logic designed, need to add TMA store for FP4 output")
|
||||
|
||||
# For now, verify that the Python quantization matches our dequantize
|
||||
def test_nvfp4_kernel():
|
||||
print("\n=== NVFP4 Kernel Quantization Test ===")
|
||||
torch.manual_seed(42)
|
||||
M, N = 4, 64
|
||||
M, N = 128, 512
|
||||
|
||||
x = torch.randn(M, N, dtype=torch.bfloat16, device='cuda')
|
||||
x_fp4, sf = quantize_activation_nvfp4(x, 1.0)
|
||||
x_deq = dequantize_nvfp4_simple(x_fp4, sf, 1.0, N)
|
||||
cos = torch.nn.functional.cosine_similarity(x.flatten().float().unsqueeze(0), x_deq.flatten().float().unsqueeze(0)).item()
|
||||
print(f" Small test cos: {cos:.6f}")
|
||||
x_fp4_ref, sf_ref = quantize_activation_nvfp4(x, 1.0)
|
||||
x_deq_ref = dequantize_nvfp4(x_fp4_ref, sf_ref, 1.0, N)
|
||||
|
||||
kernel = Nvfp4QuantizeKernel(block_size=16)
|
||||
|
||||
x_fp4_out = torch.zeros(M, N // 2, dtype=torch.uint8, device='cuda')
|
||||
sf_out = torch.zeros(M, N // 16, dtype=torch.float8_e4m3fn, device='cuda')
|
||||
|
||||
stream = cuda.CUstream(0)
|
||||
|
||||
x_bf16_cute = ct.from_dlpack(x)
|
||||
x_fp4_cute = ct.from_dlpack(x_fp4_out)
|
||||
sf_cute = ct.from_dlpack(sf_out)
|
||||
|
||||
kernel(x_bf16_cute, x_fp4_cute, sf_cute, Int32(M), Int32(N), stream)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
print(f" FP4 nonzero: {(x_fp4_out > 0).sum().item()} / {x_fp4_out.numel()}")
|
||||
print(f" SF nonzero: {(sf_out.float() > 0).sum().item()} / {sf_out.numel()}")
|
||||
print(f" FP4 sample: {x_fp4_out[0, :8]}")
|
||||
print(f" SF sample: {sf_out[0, :4].float()}")
|
||||
|
||||
x_deq_kernel = dequantize_nvfp4(x_fp4_out, sf_out, 1.0, N)
|
||||
|
||||
cos_kernel = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), x_deq_kernel.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
cos_ref = torch.nn.functional.cosine_similarity(
|
||||
x.flatten().float().unsqueeze(0), x_deq_ref.flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
|
||||
print(f" Reference Python cos: {cos_ref:.6f}")
|
||||
print(f" Kernel cos: {cos_kernel:.6f}")
|
||||
|
||||
if cos_kernel >= 0.95:
|
||||
print(f" ✅ Kernel quantization PASS (cos={cos_kernel:.4f})")
|
||||
else:
|
||||
print(f" ❌ Kernel quantization FAIL (cos={cos_kernel:.4f})")
|
||||
|
||||
fp4_match = (x_fp4_out == x_fp4_ref.view(torch.uint8)).float().mean().item()
|
||||
sf_match = (sf_out.float() == sf_ref.float()).float().mean().item()
|
||||
print(f" FP4 byte match rate: {fp4_match:.4f}")
|
||||
print(f" FP8 scale match rate: {sf_match:.4f}")
|
||||
|
||||
|
||||
def test():
|
||||
print("=== NVFP4-1.1: BF16→FP4 Quantization ===")
|
||||
print("=== NVFP4-1.1: BF4→FP4 Quantization Kernel ===")
|
||||
test_nvfp4_python()
|
||||
test_nvfp4_kernel_launch()
|
||||
test_nvfp4_kernel()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Reference in New Issue
Block a user