cleanup: move useful tests to tests/, nuke stale debug tests
Kept (moved to tests/): - test_uniform_fp4.py — proves GEMM math (72.0 = 1.5² × K) - test_b_layout.py — proves B matrix column layout - test_quick_rand.py — quick GEMM sanity check Removed (stale SF remap debug artifacts): - test_forward_map.py, test_gemm_sweep.py, test_m1_gemm.py - test_minimal_gemm.py, test_rand_gemm.py, test_sf_check.py - test_sf_remap.py, test_sf_signed.py, test_sf_layout_diag.cu
This commit is contained in:
@@ -1,84 +0,0 @@
|
||||
"""Test: verify that layout_sf(make_coord(m, k*16)) produces correct dst indices.
|
||||
If the forward mapping is wrong, this will show it."""
|
||||
import torch, sys
|
||||
sys.path.insert(0, 'src')
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
# Test 1: all-ones SF (should still give cosine 1.0)
|
||||
M, N, K = 1, 32, 32
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T; w_sf = w_sf.T
|
||||
|
||||
# Test with uniform SF
|
||||
x_sf_ones = torch.ones_like(x_sf)
|
||||
w_sf_ones = torch.ones_like(w_sf)
|
||||
|
||||
out_uni = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf_ones, w_fp4, w_sf_ones, M, N, K, alpha=1.0)
|
||||
|
||||
# Dequant reference with uniform SF
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
lo = (x_u8 & 0x0F).long(); hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_deq = ((x_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_recon = (x_deq * 1.0).to(torch.bfloat16)
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long(); whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_deq = ((w_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_recon = (w_deq * 1.0).to(torch.bfloat16)
|
||||
|
||||
ref_uni = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
cos_uni = torch.nn.functional.cosine_similarity(out_uni.float(), ref_uni.float(), dim=-1).mean().item()
|
||||
print(f"Uniform SF: cosine={cos_uni:.6f}")
|
||||
|
||||
# Test 2: try the prepack path
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import prepack_sfb
|
||||
w_sf_packed = prepack_sfb(w_sf, M, N, K)
|
||||
out_prepacked = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf_packed, M, N, K, alpha=1.0, sfb_prepacked=True)
|
||||
|
||||
# Full dequant reference
|
||||
x_recon_real = (x_deq * x_sf.to(torch.float32).repeat_interleave(16, dim=-1)).to(torch.bfloat16)
|
||||
w_recon_real = (w_deq * w_sf.to(torch.float32).repeat_interleave(16, dim=0)).to(torch.bfloat16)
|
||||
ref_real = torch.nn.functional.linear(x_recon_real, w_recon_real.T)
|
||||
|
||||
cos_pre = torch.nn.functional.cosine_similarity(out_prepacked.float(), ref_real.float(), dim=-1).mean().item()
|
||||
print(f"Prepacked SFB: cosine={cos_pre:.6f}")
|
||||
|
||||
# Test 3: without prepack (on-the-fly SFB remap)
|
||||
out_live = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
cos_live = torch.nn.functional.cosine_similarity(out_live.float(), ref_real.float(), dim=-1).mean().item()
|
||||
print(f"Live SFB: cosine={cos_live:.6f}")
|
||||
|
||||
# Test 4: N=128, K=256 (bigger dims)
|
||||
M2, N2, K2 = 1, 128, 256
|
||||
x2 = torch.randn(M2, K2, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w2 = torch.randn(K2, N2, dtype=torch.bfloat16, device=device) * 0.5
|
||||
x2_fp4, x2_sf = _quantize_to_e2m1(x2.float())
|
||||
w2_fp4, w2_sf = _quantize_to_e2m1(w2.T.float())
|
||||
w2_fp4 = w2_fp4.T; w2_sf = w2_sf.T
|
||||
|
||||
out2 = cutlass_nvfp4_blockscaled_gemm(x2_fp4, x2_sf, w2_fp4, w2_sf, M2, N2, K2, alpha=1.0)
|
||||
# Dequant ref
|
||||
x2_u8 = x2_fp4.view(torch.uint8)
|
||||
lo2 = (x2_u8 & 0x0F).long(); hi2 = ((x2_u8 >> 4) & 0x0F).long()
|
||||
x2_nib = torch.stack([lo2, hi2], dim=-1).reshape(M2, -1)
|
||||
x2_deq = ((x2_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(x2_nib & 0x07)]
|
||||
x2_recon = (x2_deq * x2_sf.to(torch.float32).repeat_interleave(16, dim=-1)).to(torch.bfloat16)
|
||||
|
||||
w2_u8 = w2_fp4.view(torch.uint8)
|
||||
w2lo = (w2_u8 & 0x0F).long(); w2hi = ((w2_u8 >> 4) & 0x0F).long()
|
||||
w2_nib = torch.stack([w2lo, w2hi], dim=-1).reshape(w2_u8.shape[0]*2, w2_u8.shape[1])
|
||||
w2_deq = ((w2_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(w2_nib & 0x07)]
|
||||
w2_recon = (w2_deq * w2_sf.to(torch.float32).repeat_interleave(16, dim=0)).to(torch.bfloat16)
|
||||
ref2 = torch.nn.functional.linear(x2_recon, w2_recon.T)
|
||||
cos2 = torch.nn.functional.cosine_similarity(out2.float(), ref2.float(), dim=-1).mean().item()
|
||||
print(f"M=1 N=128 K=256: cosine={cos2:.6f}")
|
||||
@@ -1,77 +0,0 @@
|
||||
"""Minimal test: CUTLASS NVFP4 GEMM with simple dimensions to isolate the bug.
|
||||
|
||||
Test 1: Small dimensions (M=128, N=256, K=512) — should match the original working test
|
||||
Test 2: Medium dimensions (M=4, N=1024, K=2048)
|
||||
Test 3: Real MoE dimensions (M=1, N=6144, K=7168)
|
||||
"""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, 'src')
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
def test_gemm(M, N, K, label):
|
||||
K_half = K // 2
|
||||
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
# Quantize
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
# Weight: quantize transposed to get (K_half, N) layout
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T # (K_half, N)
|
||||
w_sf = w_sf.T # (K//16, N)
|
||||
|
||||
# Dequant reference
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
lo = (x_u8 & 0x0F).long()
|
||||
hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_signs = (x_nib >> 3).float() * -2 + 1
|
||||
x_mags = _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_deq = x_signs * x_mags
|
||||
sf_exp = x_sf.to(torch.float32).repeat_interleave(16, dim=-1)
|
||||
x_recon = (x_deq * sf_exp).to(torch.bfloat16)
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long()
|
||||
whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_signs = (w_nib >> 3).float() * -2 + 1
|
||||
w_mags = _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_deq = w_signs * w_mags
|
||||
w_sf_exp = w_sf.to(torch.float32).repeat_interleave(16, dim=0)
|
||||
w_recon = (w_deq * w_sf_exp).to(torch.bfloat16)
|
||||
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).mean().item()
|
||||
mse = (nvfp4_out.float() - quant_ref.float()).pow(2).mean().item()
|
||||
print(f"{label}: M={M} N={N} K={K} cosine={cos:.6f} mse={mse:.4e} nvfp4_amax={nvfp4_out.abs().max():.2e} ref_amax={quant_ref.abs().max():.2e}")
|
||||
|
||||
# Test 1: Small (like original working test)
|
||||
test_gemm(128, 256, 512, "SMALL")
|
||||
test_gemm(128, 512, 1024, "MEDIUM")
|
||||
|
||||
# Test 2: N and K divisible by 128 (tile alignment)
|
||||
test_gemm(1, 128, 256, "TINY")
|
||||
test_gemm(1, 256, 512, "SMALL-M1")
|
||||
test_gemm(1, 1024, 2048, "MED-M1")
|
||||
|
||||
# Test 3: Real MoE dimensions
|
||||
test_gemm(1, 6144, 7168, "REAL-L1")
|
||||
test_gemm(1, 7168, 3072, "REAL-L2")
|
||||
|
||||
# Test 4: N=6144 K=7168 with M=128 (to see if M matters at these dims)
|
||||
test_gemm(128, 6144, 7168, "REAL-L1-M128")
|
||||
|
||||
# Test 5: Aligned versions
|
||||
test_gemm(1, 6144, 7168, "REAL-L1") # same, for reference
|
||||
test_gemm(1, 6144, 7168, "REAL-L1-no-alpha") # alpha=1.0 already
|
||||
101
test_m1_gemm.py
101
test_m1_gemm.py
@@ -1,101 +0,0 @@
|
||||
"""Standalone test matching real MoE dimensions: M=1, N=6144, K=7168.
|
||||
|
||||
The random test with M=128 showed cosine 1.0, but real inference with M=1
|
||||
shows cosine ≈ 0. This test uses deterministic data at M=1 to reproduce.
|
||||
"""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, 'src')
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import (
|
||||
cutlass_nvfp4_blockscaled_gemm,
|
||||
)
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
M, N, K = 1, 6144, 7168
|
||||
K_half = K // 2
|
||||
|
||||
# Create BF16 reference data
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
# Reference BF16 GEMM
|
||||
ref_out = torch.nn.functional.linear(x_bf16, w_bf16.T) # (M, N)
|
||||
print(f"BF16 ref: amax={ref_out.abs().max():.4e} mean={ref_out.mean():.4e}")
|
||||
|
||||
# Quantize to NVFP4
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float()) # (M, K_half) int8, (M, K//16) float8
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.float()) # (K, N_half) int8, (K, N//16) float8
|
||||
|
||||
# Need w in (K_half, N) layout for CUTLASS
|
||||
# w_bf16 is (K, N). Quantize gives w_fp4 (K, N//2). Need (K//2, N) = (3584, 6144)
|
||||
# Wait — the weight layout for CUTLASS B is (K_half, N) where the original is (K, N)
|
||||
# But _quantize_to_e2m1 on (K, N) gives (K, N//2) which is (7168, 3072)
|
||||
# We need (3584, 6144) = (K_half, N)
|
||||
# So we should quantize w_bf16.T instead: (N, K) → (N, K//2) → transpose to (K//2, N)
|
||||
w_t = w_bf16.T # (N, K) = (6144, 7168)
|
||||
w_fp4_t, w_sf_t = _quantize_to_e2m1(w_t.float()) # (N, K//2) = (6144, 3584)
|
||||
w_fp4_final = w_fp4_t.T # (K//2, N) = (3584, 6144)
|
||||
w_sf_final = w_sf_t.T # (K//16, N) = (448, 6144)
|
||||
|
||||
print(f"x_fp4: {x_fp4.shape} x_sf: {x_sf.shape}")
|
||||
print(f"w_fp4: {w_fp4_final.shape} w_sf: {w_sf_final.shape}")
|
||||
|
||||
# Dequantize and compute reference from quantized values
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
lo = (x_u8 & 0x0F).long()
|
||||
hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_signs = (x_nib >> 3).float() * -2 + 1
|
||||
x_mags = _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_deq = x_signs * x_mags
|
||||
sf_exp = x_sf.to(torch.float32).repeat_interleave(16, dim=-1)
|
||||
x_recon = (x_deq * sf_exp).to(torch.bfloat16)
|
||||
|
||||
w_u8 = w_fp4_final.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long()
|
||||
whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_signs = (w_nib >> 3).float() * -2 + 1
|
||||
w_mags = _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_deq = w_signs * w_mags
|
||||
w_sf_exp = w_sf_final.to(torch.float32).repeat_interleave(16, dim=0)
|
||||
w_recon = (w_deq * w_sf_exp).to(torch.bfloat16)
|
||||
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
print(f"Quant ref: amax={quant_ref.abs().max():.4e} mean={quant_ref.mean():.4e}")
|
||||
|
||||
# Run CUTLASS GEMM
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(
|
||||
x_fp4, x_sf,
|
||||
w_fp4_final, w_sf_final,
|
||||
M, N, K,
|
||||
alpha=1.0,
|
||||
)
|
||||
print(f"NVFP4 out: amax={nvfp4_out.abs().max():.4e} mean={nvfp4_out.mean():.4e}")
|
||||
|
||||
# Cosine similarity
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).mean().item()
|
||||
mse = (nvfp4_out.float() - quant_ref.float()).pow(2).mean().item()
|
||||
print(f"cosine={cos:.6f} mse={mse:.4e}")
|
||||
|
||||
# Also test with M=128
|
||||
M2 = 128
|
||||
x2 = torch.randn(M2, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
x2_fp4, x2_sf = _quantize_to_e2m1(x2.float())
|
||||
x2_u8 = x2_fp4.view(torch.uint8)
|
||||
lo2 = (x2_u8 & 0x0F).long()
|
||||
hi2 = ((x2_u8 >> 4) & 0x0F).long()
|
||||
x2_nib = torch.stack([lo2, hi2], dim=-1).reshape(M2, -1)
|
||||
x2_signs = (x2_nib >> 3).float() * -2 + 1
|
||||
x2_mags = _E2M1_MAGNITUDES.to(device)[(x2_nib & 0x07)]
|
||||
x2_deq = x2_signs * x2_mags
|
||||
sf2_exp = x2_sf.to(torch.float32).repeat_interleave(16, dim=-1)
|
||||
x2_recon = (x2_deq * sf2_exp).to(torch.bfloat16)
|
||||
qr2 = torch.nn.functional.linear(x2_recon, w_recon.T)
|
||||
nv2 = cutlass_nvfp4_blockscaled_gemm(x2_fp4, x2_sf, w_fp4_final, w_sf_final, M2, N, K, alpha=1.0)
|
||||
cos2 = torch.nn.functional.cosine_similarity(nv2.float(), qr2.float(), dim=-1).mean().item()
|
||||
print(f"M=128: cosine={cos2:.6f}")
|
||||
@@ -1,68 +0,0 @@
|
||||
"""Ultra-minimal test: 1 element output, manual verification."""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, 'src')
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
# Simplest: M=1, N=32, K=32
|
||||
M, N, K = 1, 32, 32
|
||||
|
||||
# All ones in BF16
|
||||
x_bf16 = torch.ones(M, K, dtype=torch.bfloat16, device=device)
|
||||
w_bf16 = torch.ones(K, N, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# Reference: all-ones @ all-ones = K = 32.0 for every element
|
||||
ref_out = torch.nn.functional.linear(x_bf16, w_bf16.T)
|
||||
print(f"BF16 ref (all-ones): {ref_out[0, :8].tolist()} (expected all 32.0)")
|
||||
|
||||
# Quantize
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T
|
||||
w_sf = w_sf.T
|
||||
|
||||
# Check what quantized values look like
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
print(f"x_fp4 first 8 bytes: {x_u8[0, :8].tolist()}")
|
||||
print(f"x_sf first 4: {x_sf[0, :4].to(torch.float32).tolist()}")
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
print(f"w_fp4 first 8 bytes: {w_u8[:8, 0].tolist()}")
|
||||
print(f"w_sf first 4: {w_sf[:4, 0].to(torch.float32).tolist()}")
|
||||
|
||||
# Dequant reference
|
||||
lo = (x_u8 & 0x0F).long()
|
||||
hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_signs = (x_nib >> 3).float() * -2 + 1
|
||||
x_mags = _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_deq = x_signs * x_mags
|
||||
sf_exp = x_sf.to(torch.float32).repeat_interleave(16, dim=-1)
|
||||
x_recon = (x_deq * sf_exp).to(torch.bfloat16)
|
||||
print(f"x_recon first 8: {x_recon[0, :8].tolist()}")
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long()
|
||||
whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_signs = (w_nib >> 3).float() * -2 + 1
|
||||
w_mags = _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_deq = w_signs * w_mags
|
||||
w_sf_exp = w_sf.to(torch.float32).repeat_interleave(16, dim=0)
|
||||
w_recon = (w_deq * w_sf_exp).to(torch.bfloat16)
|
||||
print(f"w_recon first 8: {w_recon[:8, 0].tolist()}")
|
||||
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
print(f"Quant ref first 8: {quant_ref[0, :8].tolist()}")
|
||||
|
||||
# CUTLASS GEMM
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
print(f"NVFP4 out first 8: {nvfp4_out[0, :8].tolist()}")
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).item()
|
||||
print(f"cosine={cos:.6f}")
|
||||
@@ -1,80 +0,0 @@
|
||||
"""Test: random data at small dimensions to check if non-uniform SF breaks it."""
|
||||
import torch
|
||||
import sys
|
||||
sys.path.insert(0, 'src')
|
||||
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
def test(M, N, K, label):
|
||||
K_half = K // 2
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T
|
||||
w_sf = w_sf.T
|
||||
|
||||
# Dequant reference
|
||||
def dequant_a(fp4, sf, M, K):
|
||||
u8 = fp4.view(torch.uint8)
|
||||
lo = (u8 & 0x0F).long()
|
||||
hi = ((u8 >> 4) & 0x0F).long()
|
||||
nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
signs = (nib >> 3).float() * -2 + 1
|
||||
mags = _E2M1_MAGNITUDES.to(device)[(nib & 0x07)]
|
||||
sf_exp = sf.to(torch.float32).repeat_interleave(16, dim=-1)
|
||||
return (signs * mags * sf_exp).to(torch.bfloat16)
|
||||
|
||||
def dequant_b(fp4, sf, K, N):
|
||||
u8 = fp4.view(torch.uint8)
|
||||
lo = (u8 & 0x0F).long()
|
||||
hi = ((u8 >> 4) & 0x0F).long()
|
||||
nib = torch.stack([lo, hi], dim=-1).reshape(u8.shape[0]*2, u8.shape[1])
|
||||
signs = (nib >> 3).float() * -2 + 1
|
||||
mags = _E2M1_MAGNITUDES.to(device)[(nib & 0x07)]
|
||||
sf_exp = sf.to(torch.float32).repeat_interleave(16, dim=0)
|
||||
return (signs * mags * sf_exp).to(torch.bfloat16)
|
||||
|
||||
x_recon = dequant_a(x_fp4, x_sf, M, K)
|
||||
w_recon = dequant_b(w_fp4, w_sf, K, N)
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).mean().item()
|
||||
mse = (nvfp4_out.float() - quant_ref.float()).pow(2).mean().item()
|
||||
print(f"{label}: M={M} N={N} K={K} cosine={cos:.6f} mse={mse:.4e}")
|
||||
|
||||
# All at N=32, K=32 (same as the working all-ones test)
|
||||
test(1, 32, 32, "RAND-TINY")
|
||||
test(4, 32, 32, "RAND-M4")
|
||||
test(128, 32, 32, "RAND-M128")
|
||||
|
||||
# Bigger
|
||||
test(1, 128, 256, "RAND-128x256")
|
||||
test(1, 256, 512, "RAND-256x512")
|
||||
test(128, 256, 512, "RAND-128x256x512")
|
||||
|
||||
# Test with alpha != 1.0
|
||||
print("\n--- alpha test ---")
|
||||
M, N, K = 1, 32, 32
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T; w_sf = w_sf.T
|
||||
|
||||
x_recon = dequant_a(x_fp4, x_sf, M, K)
|
||||
w_recon = dequant_b(w_fp4, w_sf, K, N)
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
|
||||
for alpha in [1.0, 0.5, 2.0, 1e-3, 4.6e-5]:
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=alpha)
|
||||
ref_scaled = quant_ref * alpha
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), ref_scaled.float(), dim=-1).item()
|
||||
print(f" alpha={alpha:.1e} cosine={cos:.6f}")
|
||||
@@ -1,45 +0,0 @@
|
||||
"""Check if size != cosize for small dimensions."""
|
||||
import torch, sys
|
||||
sys.path.insert(0, 'src')
|
||||
|
||||
# We need to construct the layouts to check
|
||||
# Replicate what the CU code does
|
||||
import cutlass_nvfp4_gemm._C as _C
|
||||
# Actually we can't easily call CUTE from Python.
|
||||
# Let's just test with progressively larger N until size != cosize matters.
|
||||
|
||||
# Alternative approach: test the _C.forward directly and check SF remap
|
||||
# by passing known SF values and seeing if they end up in the right places.
|
||||
|
||||
# Simpler: test with M=128, N=128, K=256 where tile padding definitely kicks in
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
# Test at dimensions where tiling definitely matters
|
||||
for M, N, K in [(128, 128, 256), (128, 256, 512), (1, 6144, 7168)]:
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T; w_sf = w_sf.T
|
||||
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
lo = (x_u8 & 0x0F).long(); hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_deq = ((x_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_recon = (x_deq * x_sf.to(torch.float32).repeat_interleave(16, dim=-1)).to(torch.bfloat16)
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long(); whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_deq = ((w_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_recon = (w_deq * w_sf.to(torch.float32).repeat_interleave(16, dim=0)).to(torch.bfloat16)
|
||||
|
||||
quant_ref = torch.nn.functional.linear(x_recon, w_recon.T)
|
||||
nvfp4_out = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
cos = torch.nn.functional.cosine_similarity(nvfp4_out.float(), quant_ref.float(), dim=-1).mean().item()
|
||||
print(f"M={M} N={N} K={K} cosine={cos:.6f}")
|
||||
@@ -1,101 +0,0 @@
|
||||
/** Diagnostic: print the SF layout coordinate mapping for verification.
|
||||
* Compile and run to see what (m, k_sf) each dst_idx maps to.
|
||||
*/
|
||||
#include <cstdio>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cute/tensor.hpp>
|
||||
#include <cutlass/float_subbyte.h>
|
||||
#include <cutlass/gemm/dispatch_policy.hpp>
|
||||
#include <cutlass/gemm/collective/collective_builder.hpp>
|
||||
#include <cutlass/epilogue/collective/collective_builder.hpp>
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/gemm_universal.hpp>
|
||||
|
||||
using namespace cute;
|
||||
|
||||
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
|
||||
using LayoutATag = cutlass::layout::RowMajor;
|
||||
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor;
|
||||
using ElementD = cutlass::bfloat16_t;
|
||||
using ElementC = float;
|
||||
using LayoutCTag = cutlass::layout::RowMajor;
|
||||
using LayoutDTag = cutlass::layout::RowMajor;
|
||||
using ElementAccumulator = float;
|
||||
using ElementCompute = float;
|
||||
using ArchTag = cutlass::arch::Sm100;
|
||||
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
|
||||
using MmaTileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
constexpr int InputSFVectorSize = 16;
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
MmaTileShape, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, LayoutCTag, 4,
|
||||
ElementD, LayoutDTag, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass,
|
||||
ElementA, LayoutATag, 32,
|
||||
ElementB, LayoutBTag, 32,
|
||||
ElementAccumulator,
|
||||
MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
|
||||
|
||||
int main() {
|
||||
// Test with M=1, N=32, K=32
|
||||
int M = 1, N = 32, K = 32;
|
||||
LayoutSFA layout_SFA = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
|
||||
|
||||
int total = cute::cosize(layout_SFA);
|
||||
int logical = cute::size(layout_SFA);
|
||||
printf("M=%d N=%d K=%d: cosize=%d size=%d\n", M, N, K, total, logical);
|
||||
|
||||
// Print shape and stride
|
||||
auto shape = layout_SFA.shape();
|
||||
auto stride = layout_SFA.stride();
|
||||
printf("Layout rank: %d\n", (int)cute::rank(shape));
|
||||
|
||||
// Print first 20 coordinate mappings
|
||||
int count = std::min(total, 20);
|
||||
for (int i = 0; i < count; i++) {
|
||||
auto coord = cute::idx2crd(i, layout_SFA.shape(), layout_SFA.stride());
|
||||
auto flat = cute::flatten(coord);
|
||||
constexpr int R = cute::rank_v<decltype(flat)>;
|
||||
printf("dst[%d]: rank=%d", i, R);
|
||||
if constexpr (R >= 1) printf(" f0=%d", (int)cute::get<0>(flat));
|
||||
if constexpr (R >= 2) printf(" f1=%d", (int)cute::get<1>(flat));
|
||||
if constexpr (R >= 3) printf(" f2=%d", (int)cute::get<2>(flat));
|
||||
if constexpr (R >= 4) printf(" f3=%d", (int)cute::get<3>(flat));
|
||||
if constexpr (R >= 5) printf(" f4=%d", (int)cute::get<4>(flat));
|
||||
if constexpr (R >= 6) printf(" f5=%d", (int)cute::get<5>(flat));
|
||||
if constexpr (R >= 7) printf(" f6=%d", (int)cute::get<6>(flat));
|
||||
if constexpr (R >= 8) printf(" f7=%d", (int)cute::get<7>(flat));
|
||||
|
||||
// Compute m and k_sf using the current formula
|
||||
int m = 0, k_sf = 0;
|
||||
if constexpr (R == 8) {
|
||||
m = cute::get<0>(flat) + cute::get<1>(flat) * 32 + cute::get<2>(flat) * 128;
|
||||
k_sf = cute::get<4>(flat) + cute::get<5>(flat) * 4;
|
||||
}
|
||||
printf(" -> m=%d k_sf=%d\n", m, k_sf);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -1,75 +0,0 @@
|
||||
"""Verify the SF remap by comparing CUTLASS output with and without SF remap.
|
||||
|
||||
Strategy:
|
||||
1. Run GEMM with identity SF (all 1.0) — both A and B
|
||||
2. Run GEMM with a single non-1.0 SF value — see if it affects the right output elements
|
||||
3. This tells us if the remap is placing SF values correctly
|
||||
|
||||
Actually, simpler: run GEMM with prepack_sfb=False (remap on the fly) and
|
||||
prepack_sfb=True (pre-remapped), compare. If they differ, the remap is wrong.
|
||||
"""
|
||||
import torch, sys
|
||||
sys.path.insert(0, 'src')
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import (
|
||||
cutlass_nvfp4_blockscaled_gemm, prepack_sfb
|
||||
)
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1, _E2M1_MAGNITUDES
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
|
||||
M, N, K = 1, 32, 32
|
||||
x_bf16 = torch.randn(M, K, dtype=torch.bfloat16, device=device) * 2.0
|
||||
w_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device) * 0.5
|
||||
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x_bf16.float())
|
||||
w_fp4, w_sf = _quantize_to_e2m1(w_bf16.T.float())
|
||||
w_fp4 = w_fp4.T; w_sf = w_sf.T
|
||||
|
||||
# Test 1: with remap (sfb_prepacked=False)
|
||||
out_remap = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf, M, N, K, alpha=1.0, sfb_prepacked=False)
|
||||
|
||||
# Test 2: with prepacked SFB
|
||||
w_sf_packed = prepack_sfb(w_sf, M, N, K)
|
||||
out_prepacked = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf_packed, M, N, K, alpha=1.0, sfb_prepacked=True)
|
||||
|
||||
print(f"Remap output first 8: {out_remap[0,:8].tolist()}")
|
||||
print(f"Prepacked output first 8: {out_prepacked[0,:8].tolist()}")
|
||||
print(f"Match: {torch.allclose(out_remap, out_prepacked, atol=0.01)}")
|
||||
diff = (out_remap - out_prepacked).abs().max().item()
|
||||
print(f"Max diff: {diff:.4e}")
|
||||
|
||||
# Test 3: uniform SF — should match perfectly
|
||||
x_sf_ones = torch.ones_like(x_sf)
|
||||
w_sf_ones = torch.ones_like(w_sf)
|
||||
out_uni_remap = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf_ones, w_fp4, w_sf_ones, M, N, K, alpha=1.0, sfb_prepacked=False)
|
||||
out_uni_pre = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf_ones, w_fp4, prepack_sfb(w_sf_ones, M, N, K), M, N, K, alpha=1.0, sfb_prepacked=True)
|
||||
print(f"\nUniform SF remap vs prepacked: {torch.allclose(out_uni_remap, out_uni_pre, atol=0.01)}")
|
||||
|
||||
# Test 4: SFA remap — try with all-1.0 SFA and actual SFB, vs actual SFA and all-1.0 SFB
|
||||
# This isolates which remap (SFA or SFB) is broken
|
||||
out_real_sfa = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf, w_fp4, w_sf_ones, M, N, K, alpha=1.0)
|
||||
out_real_sfb = cutlass_nvfp4_blockscaled_gemm(x_fp4, x_sf_ones, w_fp4, w_sf, M, N, K, alpha=1.0)
|
||||
|
||||
# Compute BF16 references
|
||||
x_u8 = x_fp4.view(torch.uint8)
|
||||
lo = (x_u8 & 0x0F).long(); hi = ((x_u8 >> 4) & 0x0F).long()
|
||||
x_nib = torch.stack([lo, hi], dim=-1).reshape(M, -1)
|
||||
x_deq = ((x_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(x_nib & 0x07)]
|
||||
x_recon = (x_deq * x_sf.to(torch.float32).repeat_interleave(16, dim=-1)).to(torch.bfloat16)
|
||||
x_recon_ones = (x_deq * 1.0).to(torch.bfloat16) # uniform SF
|
||||
|
||||
w_u8 = w_fp4.view(torch.uint8)
|
||||
wlo = (w_u8 & 0x0F).long(); whi = ((w_u8 >> 4) & 0x0F).long()
|
||||
w_nib = torch.stack([wlo, whi], dim=-1).reshape(w_u8.shape[0]*2, w_u8.shape[1])
|
||||
w_deq = ((w_nib >> 3).float() * -2 + 1) * _E2M1_MAGNITUDES.to(device)[(w_nib & 0x07)]
|
||||
w_recon = (w_deq * w_sf.to(torch.float32).repeat_interleave(16, dim=0)).to(torch.bfloat16)
|
||||
w_recon_ones = (w_deq * 1.0).to(torch.bfloat16)
|
||||
|
||||
ref_real_sfa = torch.nn.functional.linear(x_recon, w_recon_ones.T)
|
||||
ref_real_sfb = torch.nn.functional.linear(x_recon_ones, w_recon.T)
|
||||
|
||||
cos_sfa = torch.nn.functional.cosine_similarity(out_real_sfa.float(), ref_real_sfa.float(), dim=-1).mean().item()
|
||||
cos_sfb = torch.nn.functional.cosine_similarity(out_real_sfb.float(), ref_real_sfb.float(), dim=-1).mean().item()
|
||||
print(f"\nSFA remap cosine (real SFA, uniform SFB): {cos_sfa:.6f}")
|
||||
print(f"SFB remap cosine (uniform SFA, real SFB): {cos_sfb:.6f}")
|
||||
@@ -1,45 +0,0 @@
|
||||
"""Check if float8_e4m3fn (signed) vs float_ue4m3 (unsigned) matters.
|
||||
In the CUTLASS kernel, SF is float_ue4m3 (unsigned E4M3).
|
||||
In our Python reference, we use .to(torch.float32) which interprets float8_e4m3fn (signed).
|
||||
If the sign bit is set, signed and unsigned give different values.
|
||||
"""
|
||||
import torch
|
||||
device = "cuda"
|
||||
|
||||
# Create some float8 values and compare signed vs unsigned interpretation
|
||||
vals = torch.tensor([0x00, 0x3F, 0x7F, 0x80, 0xBF, 0xFF], dtype=torch.uint8, device=device)
|
||||
|
||||
# Signed interpretation (float8_e4m3fn)
|
||||
signed = vals.view(torch.float8_e4m3fn).to(torch.float32)
|
||||
print("Signed (float8_e4m3fn):", signed.tolist())
|
||||
|
||||
# Unsigned interpretation (float8_e4m3fnuz — unsigned zero)
|
||||
# Actually, let's check if there IS an unsigned float8 type in PyTorch
|
||||
print("Has float8_e4m3fnuz:", hasattr(torch, 'float8_e4m3fnuz'))
|
||||
|
||||
# The key question: are SF values always positive?
|
||||
# UE4M3 means the sign bit is NOT used — all values are positive.
|
||||
# But if we read a UE4M3 byte as signed E4M3, bytes with bit 7 set
|
||||
# would be interpreted as negative.
|
||||
# Let's check: for valid UE4M3 values, is bit 7 ever set?
|
||||
# E4M3 range: 0 to 448. The encoding uses the sign bit for actual sign.
|
||||
# UE4M3: the sign bit is always 0 (positive only, range 0 to 448).
|
||||
# So reading UE4M3 as signed E4M3 should give the same result
|
||||
# as long as the sign bit is 0.
|
||||
|
||||
# Check our actual SF data
|
||||
from nvfp4_megamoe_kernel.nvfp4_mega_moe import _quantize_to_e2m1
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 32, device=device) * 2.0
|
||||
x_fp4, x_sf = _quantize_to_e2m1(x.float())
|
||||
sf_bytes = x_sf.view(torch.uint8)
|
||||
print(f"\nSF bytes: {sf_bytes.flatten()[:16].tolist()}")
|
||||
print(f"Any byte with bit 7 set (>= 128): {(sf_bytes >= 128).any().item()}")
|
||||
print(f"SF as signed float: {x_sf.to(torch.float32).flatten()[:8].tolist()}")
|
||||
|
||||
# Check: does CUTLASS treat SF as signed or unsigned?
|
||||
# The C++ type is cutlass::float_ue4m3_t
|
||||
# In the CU file we use: const cutlass::float_ue4m3_t* src
|
||||
# But PyTorch passes float8_e4m3fn (signed)
|
||||
# These have the same bit pattern for positive values
|
||||
# but DIFFERENT bit patterns for values where the sign bit is set
|
||||
Reference in New Issue
Block a user