125 lines
4.8 KiB
Python
125 lines
4.8 KiB
Python
"""NVFP4-0: Verify Blackwell FP4 primitives are correct.
|
|
|
|
Quick diagnostic — prints sf_dtype, sf_vec_size, MMA kind, TMA element type.
|
|
Run on B200 only.
|
|
"""
|
|
import torch
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
from cutlass.cute.nvgpu import tcgen05
|
|
from cutlass import Float32, BFloat16, Float8E4M3FN, Float8E8M0FNU, Float4E2M1FN
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
|
|
|
|
|
def test_nvfp4_primitives():
|
|
print("=" * 60)
|
|
print("NVFP4-0.1: sf_dtype and sf_vec_size")
|
|
print("=" * 60)
|
|
|
|
from dsv4.ops.quantize import SF_VEC_SIZE
|
|
print(f" quantize.py SF_VEC_SIZE = {SF_VEC_SIZE}")
|
|
assert SF_VEC_SIZE == 16, f"SF_VEC_SIZE should be 16 for NVFP4, got {SF_VEC_SIZE}"
|
|
print(f" ✅ SF_VEC_SIZE = 16 (NVFP4 correct)")
|
|
|
|
# Construct a BlockScaledGEMM with E4M3 scales (NVFP4)
|
|
M, N, K = 128, 256, 512
|
|
a_fp4 = torch.randint(0, 256, (M, K // 2), device='cuda', dtype=torch.uint8).view(torch.float4_e2m1fn_x2)
|
|
b_fp4 = torch.randint(0, 256, (K // 2, N), device='cuda', dtype=torch.uint8).view(torch.float4_e2m1fn_x2)
|
|
a_sf = torch.randint(0, 256, (M, K // 16), device='cuda', dtype=torch.uint8).view(torch.float8_e4m3fn)
|
|
b_sf = torch.randint(0, 256, (N, K // 16), device='cuda', dtype=torch.uint8).view(torch.float8_e4m3fn)
|
|
|
|
from dsv4.kernels.gemm.dense import Sm100BlockScaledPersistentDenseGemmKernel
|
|
try:
|
|
gemm = Sm100BlockScaledPersistentDenseGemmKernel(
|
|
a_ptr=a_fp4,
|
|
b_ptr=b_fp4,
|
|
sfa_ptr=a_sf,
|
|
sfb_ptr=b_sf,
|
|
sf_vec_size=16,
|
|
)
|
|
print(f" DenseGEMM.sf_vec_size = {gemm.sf_vec_size}")
|
|
print(f" DenseGEMM.sf_dtype = {gemm.sf_dtype}")
|
|
if gemm.sf_dtype == Float8E4M3FN:
|
|
print(f" ✅ sf_dtype is Float8E4M3FN (NVFP4 correct)")
|
|
elif gemm.sf_dtype == Float8E8M0FNU:
|
|
print(f" ⚠️ sf_dtype is Float8E8M0FNU — this is MXFP4 scale format, NOT NVFP4!")
|
|
else:
|
|
print(f" ❌ sf_dtype is {gemm.sf_dtype} — unexpected!")
|
|
except Exception as e:
|
|
print(f" DenseGEMM construction failed: {e}")
|
|
|
|
# Also try with E8M0 scales (MXFP4)
|
|
a_sf_e8m0 = torch.randint(0, 256, (M, K // 32), device='cuda', dtype=torch.uint8).view(torch.float8_e8m0fnu)
|
|
b_sf_e8m0 = torch.randint(0, 256, (N, K // 32), device='cuda', dtype=torch.uint8).view(torch.float8_e8m0fnu)
|
|
try:
|
|
gemm_e8m0 = Sm100BlockScaledPersistentDenseGemmKernel(
|
|
a_ptr=a_fp4,
|
|
b_ptr=b_fp4,
|
|
sfa_ptr=a_sf_e8m0,
|
|
sfb_ptr=b_sf_e8m0,
|
|
sf_vec_size=32,
|
|
)
|
|
print(f" DenseGEMM (E8M0/vs=32).sf_dtype = {gemm_e8m0.sf_dtype}")
|
|
print(f" DenseGEMM (E8M0/vs=32).sf_vec_size = {gemm_e8m0.sf_vec_size}")
|
|
except Exception as e:
|
|
print(f" DenseGEMM (E8M0) failed: {e}")
|
|
|
|
# Check the grouped kernel too
|
|
from dsv4.kernels.gemm.grouped import ScaledGroupedGemmKernel
|
|
try:
|
|
grp = ScaledGroupedGemmKernel(
|
|
a_ptr=a_fp4,
|
|
b_ptr=b_fp4,
|
|
sfa_ptr=a_sf,
|
|
sfb_ptr=b_sf,
|
|
sf_vec_size=16,
|
|
)
|
|
print(f" GroupedGEMM.sf_vec_size = {grp.sf_vec_size}")
|
|
print(f" GroupedGEMM.sf_dtype = {grp.sf_dtype}")
|
|
except Exception as e:
|
|
print(f" GroupedGEMM construction failed: {e}")
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("NVFP4-0.3: FP4 TMA element type in quantize.py")
|
|
print("=" * 60)
|
|
|
|
from dsv4.ops.quantize import quantize_tensor_nvfp4
|
|
x = torch.randn(4, 512, device='cuda', dtype=torch.bfloat16)
|
|
x_fp4, x_sf = quantize_tensor_nvfp4(x)
|
|
print(f" Input dtype: {x.dtype}")
|
|
print(f" FP4 output dtype: {x_fp4.dtype}")
|
|
print(f" SF output dtype: {x_sf.dtype}")
|
|
print(f" FP4 shape: {x_fp4.shape} (expected: [4, 256])")
|
|
print(f" SF shape: {x_sf.shape} (expected: [4, 32])")
|
|
if x_fp4.dtype == torch.float4_e2m1fn_x2:
|
|
print(f" ✅ FP4 tensor is float4_e2m1fn_x2 — correct for TMA")
|
|
else:
|
|
print(f" ❌ FP4 tensor dtype is {x_fp4.dtype} — should be float4_e2m1fn_x2!")
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("NVFP4-0.4: MMA kind verification")
|
|
print("=" * 60)
|
|
|
|
try:
|
|
a_major = cutlass.utils.LayoutEnum.ROW_MAJOR.mma_major_mode()
|
|
b_major = cutlass.utils.LayoutEnum.COLUMN_MAJOR.mma_major_mode()
|
|
mma = cutlass.utils.sm100.make_trivial_tiled_mma(
|
|
Float4E2M1FN, Float4E2M1FN, a_major, b_major, Float32,
|
|
tcgen05.CtaGroup.ONE, (128, 256),
|
|
tcgen05.OperandSource.SMEM,
|
|
)
|
|
print(f" FP4 MMA shape_mnk = {mma.shape_mnk}")
|
|
print(f" ✅ FP4 MMA construction succeeded")
|
|
except Exception as e:
|
|
print(f" FP4 MMA construction failed: {e}")
|
|
|
|
print()
|
|
print("DONE — NVFP4-0 verification complete")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_nvfp4_primitives()
|