From 78bebff736a2a947c33d7433a9bcb502103e038f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 05:55:45 +0000 Subject: [PATCH] test: standalone CuTeDSL GEMM diagnostic --- tests/test_cutedsl_gemm.py | 225 +++++++++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 tests/test_cutedsl_gemm.py diff --git a/tests/test_cutedsl_gemm.py b/tests/test_cutedsl_gemm.py new file mode 100644 index 00000000..b2e43d47 --- /dev/null +++ b/tests/test_cutedsl_gemm.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +""" +Standalone CuTeDSL NVFP4 GEMM test. + +Tests the ScaledGroupedGemmKernel directly with simple tensors. +Goal: confirm the kernel produces nonzero output and measure cosine vs BF16. + +Usage (on B200, in the test venv): + source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate + python3 test_cutedsl_gemm.py +""" +import torch +import sys +import os + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) +sys.path.insert(0, os.path.join(REPO_ROOT, "src")) + +from cutedsl.bridge import ( + quantize_to_nvfp4, + quantize_weight_to_nvfp4, + make_b_k_major, + assemble_scales_2d_side, + assemble_scales_3d_side, + run_nvfp4_grouped_gemm, +) + +DEVICE = "cuda:0" + +def test_small_gemm(): + """Test 128x128 GEMM: A(128,128) x B(1,128,128) -> C(128,128)""" + print("=" * 60) + print("Test 1: Small 128x128 GEMM") + print("=" * 60) + + torch.manual_seed(42) + M, K, N = 128, 128, 128 + A = torch.randn(M, K, dtype=torch.bfloat16, device=DEVICE) + B = torch.randn(K, N, dtype=torch.bfloat16, device=DEVICE) + C_ref = A @ B + + A_fp4, A_sf, A_gs = quantize_to_nvfp4(A) + B_fp4, B_sf, B_gs = quantize_weight_to_nvfp4(B) + + B_mat = make_b_k_major(B_fp4.unsqueeze(0)) + scale_a = assemble_scales_2d_side([A_sf]) + scale_b = assemble_scales_3d_side([B_sf]) + gs_a = torch.tensor([A_gs], dtype=torch.float32, device=DEVICE) + gs_b = torch.tensor([B_gs], dtype=torch.float32, device=DEVICE) + offsets = torch.tensor([0, M], dtype=torch.int32, device=DEVICE) + + print(f" A_fp4: {A_fp4.shape}, B_mat: {B_mat.shape}") + print(f" scale_a: {scale_a.shape}, scale_b: {scale_b.shape}") + print(f" gs_a: {A_gs:.8f}, gs_b: {B_gs:.8f}") + + C_out = run_nvfp4_grouped_gemm( + mat_a=A_fp4, mat_b=B_mat, + scale_a=scale_a, scale_b=scale_b, + expert_offsets=offsets, + global_scale_a=gs_a, global_scale_b=gs_b, + ) + + nonzero = (C_out.abs() > 1e-6).sum().item() + print(f" Output: shape={C_out.shape}, nonzero={nonzero}/{C_out.numel()}") + if nonzero > 0: + cos = torch.nn.functional.cosine_similarity( + C_out.float().unsqueeze(0), C_ref.float().unsqueeze(0) + ).item() + print(f" ✅ Cosine vs BF16: {cos:.6f}") + print(f" amax: {C_out.abs().max():.4f} vs ref {C_ref.abs().max():.4f}") + else: + print(f" ❌ ALL ZEROS - kernel not producing output") + return C_out, C_ref + + +def test_direct_kernel(): + """Bypass bridge.py, call the kernel directly to debug.""" + print("\n" + "=" * 60) + print("Test 2: Direct kernel invocation (bypass bridge)") + print("=" * 60) + + import cutlass + import cutlass.cute as cute + import cutlass.torch as cutlass_torch + import cutlass.utils as utils + import cuda.bindings.driver as cuda_driver + from cutedsl.kernel.moe.torch_scaled_grouped_mm import ( + ScaledGroupedGemmKernel, + pad_and_swizzle_single, + assemble_raw_scales_2d3d_2d_side, + assemble_raw_scales_2d3d_3d_side, + ceil_div, + ) + + torch.manual_seed(42) + M, K, N = 128, 128, 128 + A = torch.randn(M, K, dtype=torch.bfloat16, device=DEVICE) + B = torch.randn(K, N, dtype=torch.bfloat16, device=DEVICE) + + A_fp4, A_sf, A_gs = quantize_to_nvfp4(A) + B_fp4, B_sf, B_gs = quantize_weight_to_nvfp4(B) + + # Build B with K-major layout + B_mat = make_b_k_major(B_fp4.unsqueeze(0)) + + # Build scales using raw assembly (same as bridge but more explicit) + # 2D side: activation scales + scale_a = assemble_scales_2d_side([A_sf]) + # 3D side: weight scales + scale_b = assemble_scales_3d_side([B_sf]) + + # Global scales + gsa = torch.tensor([A_gs], dtype=torch.float32, device=DEVICE) + gsb = torch.tensor([B_gs], dtype=torch.float32, device=DEVICE) + + # Output (pre-fill with sentinel to detect writes) + out = torch.full((M, N), 42.0, dtype=torch.bfloat16, device=DEVICE) + offsets = torch.tensor([0, M], dtype=torch.int32, device=DEVICE) + + # Create kernel + kernel = ScaledGroupedGemmKernel( + scenario="2Dx3D", + sf_vec_size=16, + accumulate_on_output=False, + separate_tensormap_init=True, + consistent_token_padding=False, + mma_tiler_mnk=(128, 128, 256), + cluster_shape_mnk=(1, 1, 1), + ) + + def to_cute(t): + ct = cutlass_torch.from_dlpack(t) + return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) + + ws_size = kernel.get_workspace_size(1) + workspace = torch.full((ws_size,), 255, dtype=torch.uint8, device=DEVICE) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + mac = utils.HardwareInfo().get_max_active_clusters(1) + + print(f" Compiling kernel...") + compiled = cute.compile( + kernel, + to_cute(A_fp4), to_cute(B_mat), + to_cute(scale_a), to_cute(scale_b), + to_cute(out), to_cute(offsets), to_cute(workspace), + mac, stream, + global_scale_a=to_cute(gsa), global_scale_b=to_cute(gsb), + ) + print(f" Compilation done. Running...") + + # Run + compiled( + to_cute(A_fp4), to_cute(B_mat), + to_cute(scale_a), to_cute(scale_b), + to_cute(out), to_cute(offsets), to_cute(workspace), + stream, + global_scale_a=to_cute(gsa), global_scale_b=to_cute(gsb), + ) + torch.cuda.synchronize() + + still_42 = (out == 42.0).sum().item() + changed = out.numel() - still_42 + print(f" Output: still_sentinel={still_42}, changed={changed}/{out.numel()}") + if changed > 0: + nonzero = (out.abs() > 1e-6).sum().item() + print(f" Nonzero values: {nonzero}/{out.numel()}") + if nonzero > 0: + C_ref = A @ B + cos = torch.nn.functional.cosine_similarity( + out.float().unsqueeze(0), C_ref.float().unsqueeze(0) + ).item() + print(f" ✅ Cosine vs BF16: {cos:.6f}") + else: + print(f" ❌ Kernel did NOT write to output tensor at all") + + return out + + +def test_larger_gemm(): + """Test with dimensions matching the actual MoE weights.""" + print("\n" + "=" * 60) + print("Test 3: MoE-sized GEMM (7168 x 6144)") + print("=" * 60) + + torch.manual_seed(42) + M, K, N = 4, 7168, 6144 + A = torch.randn(M, K, dtype=torch.bfloat16, device=DEVICE) + B = torch.randn(K, N, dtype=torch.bfloat16, device=DEVICE) + C_ref = A @ B + + A_fp4, A_sf, A_gs = quantize_to_nvfp4(A) + B_fp4, B_sf, B_gs = quantize_weight_to_nvfp4(B) + + B_mat = make_b_k_major(B_fp4.unsqueeze(0)) + scale_a = assemble_scales_2d_side([A_sf]) + scale_b = assemble_scales_3d_side([B_sf]) + gs_a = torch.tensor([A_gs], dtype=torch.float32, device=DEVICE) + gs_b = torch.tensor([B_gs], dtype=torch.float32, device=DEVICE) + offsets = torch.tensor([0, M], dtype=torch.int32, device=DEVICE) + + C_out = run_nvfp4_grouped_gemm( + mat_a=A_fp4, mat_b=B_mat, + scale_a=scale_a, scale_b=scale_b, + expert_offsets=offsets, + global_scale_a=gs_a, global_scale_b=gs_b, + ) + + nonzero = (C_out.abs() > 1e-6).sum().item() + print(f" Output: shape={C_out.shape}, nonzero={nonzero}/{C_out.numel()}") + if nonzero > 0: + cos = torch.nn.functional.cosine_similarity( + C_out.float().unsqueeze(0), C_ref.float().unsqueeze(0) + ).item() + print(f" ✅ Cosine vs BF16: {cos:.6f}") + else: + print(f" ❌ ALL ZEROS") + return C_out, C_ref + + +if __name__ == "__main__": + test_small_gemm() + test_direct_kernel() + # test_larger_gemm() # uncomment after small test works + print("\nDone.")