From 303b6a89938f435e2d062379155cacdb5acc4f9f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 02:14:37 +0000 Subject: [PATCH] cleanup: move useful tests to tests/, nuke stale debug tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- test_forward_map.py | 84 --------------- test_gemm_sweep.py | 77 ------------- test_m1_gemm.py | 101 ------------------ test_minimal_gemm.py | 68 ------------ test_rand_gemm.py | 80 -------------- test_sf_check.py | 45 -------- test_sf_layout_diag.cu | 101 ------------------ test_sf_remap.py | 75 ------------- test_sf_signed.py | 45 -------- test_b_layout.py => tests/test_b_layout.py | 0 .../test_quick_rand.py | 0 .../test_uniform_fp4.py | 0 12 files changed, 676 deletions(-) delete mode 100644 test_forward_map.py delete mode 100644 test_gemm_sweep.py delete mode 100644 test_m1_gemm.py delete mode 100644 test_minimal_gemm.py delete mode 100644 test_rand_gemm.py delete mode 100644 test_sf_check.py delete mode 100644 test_sf_layout_diag.cu delete mode 100644 test_sf_remap.py delete mode 100644 test_sf_signed.py rename test_b_layout.py => tests/test_b_layout.py (100%) rename test_quick_rand.py => tests/test_quick_rand.py (100%) rename test_uniform_fp4.py => tests/test_uniform_fp4.py (100%) diff --git a/test_forward_map.py b/test_forward_map.py deleted file mode 100644 index c344e5d4..00000000 --- a/test_forward_map.py +++ /dev/null @@ -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}") diff --git a/test_gemm_sweep.py b/test_gemm_sweep.py deleted file mode 100644 index 9680627e..00000000 --- a/test_gemm_sweep.py +++ /dev/null @@ -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 diff --git a/test_m1_gemm.py b/test_m1_gemm.py deleted file mode 100644 index 1375595f..00000000 --- a/test_m1_gemm.py +++ /dev/null @@ -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}") diff --git a/test_minimal_gemm.py b/test_minimal_gemm.py deleted file mode 100644 index a001dda5..00000000 --- a/test_minimal_gemm.py +++ /dev/null @@ -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}") diff --git a/test_rand_gemm.py b/test_rand_gemm.py deleted file mode 100644 index a5f60ee8..00000000 --- a/test_rand_gemm.py +++ /dev/null @@ -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}") diff --git a/test_sf_check.py b/test_sf_check.py deleted file mode 100644 index 0611ace3..00000000 --- a/test_sf_check.py +++ /dev/null @@ -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}") diff --git a/test_sf_layout_diag.cu b/test_sf_layout_diag.cu deleted file mode 100644 index b7ec0a57..00000000 --- a/test_sf_layout_diag.cu +++ /dev/null @@ -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 -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace cute; - -using ElementA = cutlass::nv_float4_t; -using LayoutATag = cutlass::layout::RowMajor; -using ElementB = cutlass::nv_float4_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(sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::collective::KernelScheduleAuto ->::CollectiveOp; - -using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, - CollectiveMainloop, - CollectiveEpilogue, - void>; - -using Gemm = cutlass::gemm::device::GemmUniversalAdapter; -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; - 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; -} diff --git a/test_sf_remap.py b/test_sf_remap.py deleted file mode 100644 index bcf70b81..00000000 --- a/test_sf_remap.py +++ /dev/null @@ -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}") diff --git a/test_sf_signed.py b/test_sf_signed.py deleted file mode 100644 index c25f3db3..00000000 --- a/test_sf_signed.py +++ /dev/null @@ -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 diff --git a/test_b_layout.py b/tests/test_b_layout.py similarity index 100% rename from test_b_layout.py rename to tests/test_b_layout.py diff --git a/test_quick_rand.py b/tests/test_quick_rand.py similarity index 100% rename from test_quick_rand.py rename to tests/test_quick_rand.py diff --git a/test_uniform_fp4.py b/tests/test_uniform_fp4.py similarity index 100% rename from test_uniform_fp4.py rename to tests/test_uniform_fp4.py