test: standalone CuTeDSL GEMM diagnostic
This commit is contained in:
225
tests/test_cutedsl_gemm.py
Normal file
225
tests/test_cutedsl_gemm.py
Normal file
@@ -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.")
|
||||
Reference in New Issue
Block a user