From f9330a17770b97e0540cc9e6319977c99091de9d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 18:47:26 +0000 Subject: [PATCH] test: standalone M=1 GEMM test with deterministic data --- test_m1_gemm.py | 101 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 test_m1_gemm.py diff --git a/test_m1_gemm.py b/test_m1_gemm.py new file mode 100644 index 00000000..1457ac28 --- /dev/null +++ b/test_m1_gemm.py @@ -0,0 +1,101 @@ +"""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().unsqueeze(0), quant_ref.float().unsqueeze(0)).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().unsqueeze(0), qr2.float().unsqueeze(0)).item() +print(f"M=128: cosine={cos2:.6f}")