test: fix gemm 1-group test params

test_gemm_1group.py

@@ -1,56 +1,47 @@
 #!/usr/bin/env python3
-"""Test: run_nvfp4_grouped_gemm with num_groups=1 on different GPUs."""
+"""Test: run_nvfp4_grouped_gemm with 1 expert on different GPUs."""
 import torch
 from dsv4.ops.gemm_runner import run_nvfp4_grouped_gemm
-from dsv4.ops.quantize import quantize_nvfp4_gpu
+from dsv4.ops.quantize import quantize_nvfp4_gpu, quantize_weight_to_nvfp4
 from dsv4.ops.layouts import make_b_k_major, assemble_scales_3d_side
 
 torch.manual_seed(42)
 
-# Create a simple BF16 weight and quantize
 M, N, K = 1, 3072, 7168
 
 for gpu in [0, 1]:
     torch.cuda.set_device(gpu)
     dev = f"cuda:{gpu}"
     
-    # Weight
     w = torch.randn(N, K, dtype=torch.bfloat16, device=dev)
-    from dsv4.ops.quantize import quantize_weight_to_nvfp4
     w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(w)
-    gsb = torch.tensor([1.0], dtype=torch.float32, device=dev)  # simplified gsb
     
-    # K-major layout
+    # K-major layout (1 expert)
     w_km = make_b_k_major(w_fp4.unsqueeze(0))  # (1, K_sf, N)
     w_sf_3d = assemble_scales_3d_side(w_sf.unsqueeze(0))  # (1, K_sf_padded, N)
     
     # Activation
-    x = torch.randn(M, K, dtype=torch.bfloat16, device=dev)
+    x = torch.randn(128, K, dtype=torch.bfloat16, device=dev)  # padded to 128
     gsa = 1.0 / (6.0 * 448.0)
     x_fp4, x_sf = quantize_nvfp4_gpu(x, gsa)
     
-    # Expert offsets
-    padded_rows = 128
-    expert_offsets = torch.tensor([padded_rows], dtype=torch.int32, device=dev)
-    
-    # Output
-    out = torch.zeros(padded_rows, N, dtype=torch.bfloat16, device=dev)
+    # Expert offsets (1 expert, 128 rows)
+    expert_offsets = torch.tensor([128], dtype=torch.int32, device=dev)
     
     # Global scales
     gsa_buf = torch.tensor([gsa], dtype=torch.float32, device=dev)
+    gsb = torch.tensor([1.0], dtype=torch.float32, device=dev)
     
-    # Run GEMM with 1 group
-    run_nvfp4_grouped_gemm(
-        mat_a=x_fp4[:padded_rows],
-        scale_a=x_sf[:padded_rows],
+    # Run
+    out = run_nvfp4_grouped_gemm(
+        mat_a=x_fp4,
+        scale_a=x_sf,
         mat_b=w_km,
         scale_b=w_sf_3d,
         expert_offsets=expert_offsets,
         global_scale_a=gsa_buf,
         global_scale_b=gsb,
-        out=out,
-        num_groups=1,
     )
     
     has_nan = torch.isnan(out[:M]).any().item()
-    print(f"GPU {gpu}: |out|={out[:M].abs().max().item() if not has_nan else 'NaN'} has_nan={has_nan} shape={out[:M].shape}")
+    print(f"GPU {gpu}: |out|={out[:M].abs().max().item() if not has_nan else 'NaN'} has_nan={has_nan} shape={out.shape}")