Fix fused router test: use quantize_weight_to_nvfp4 (correct function name)

tests/unit/test_fused_router.py

@@ -34,9 +34,7 @@ def _reference_topk(
     act = _sqrt_softplus(logits)
     scores = act + e_bias.unsqueeze(0)
     topk_vals, topk_ids = scores.topk(top_k, dim=1)
-    # Gather activations for the selected experts
     selected_act = act.gather(1, topk_ids)
-    # Renormalize
     act_sum = selected_act.sum(dim=1, keepdim=True)
     weights = selected_act / act_sum * routed_scaling_factor
     return weights, topk_ids
@@ -56,21 +54,18 @@ def test_reference_router():
     W_gate = torch.randn(K, N, dtype=torch.bfloat16, device=device)
     e_bias = torch.randn(N, dtype=torch.float32, device=device)
 
-    # Reference: BF16 linear + activation_topk
     logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float())
     from dsv4.kernels.router._activation_topk import run_fused_activation_topk
     out_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
     out_ids = torch.empty(M, top_k, dtype=torch.int32, device=device)
     run_fused_activation_topk(logits, e_bias, routed_scaling_factor, top_k, out_w, out_ids)
 
-    # Also compute pure PyTorch reference
     ref_w, ref_ids = _reference_topk(logits, e_bias, routed_scaling_factor, top_k)
 
     print(f"Reference router test (M={M}, K={K}, N={N}, top_k={top_k}):")
     print(f"  Top-k IDs (row 0): {out_ids[0].tolist()}")
     print(f"  Top-k weights (row 0): {[f'{w:.4f}' for w in out_w[0].tolist()]}")
 
-    # Verify weights sum to routed_scaling_factor
     w_sum = out_w.sum(dim=1)
     for row in range(M):
         diff = abs(w_sum[row].item() - routed_scaling_factor)
@@ -88,7 +83,6 @@ def test_reference_router():
     assert (out_w >= 0).all(), "Negative weights"
     print(f"  All weights non-negative OK")
 
-    # Cross-check with pure PyTorch reference
     ids_match = (out_ids == ref_ids).all().item()
     if ids_match:
         print(f"  IDs match PyTorch reference OK")
@@ -99,6 +93,25 @@ def test_reference_router():
     print("Reference router test PASSED")
 
 
+def _build_nvfp4_linear(K, N, W_gate_bf16, device):
+    """Build Nvfp4Linear with NVFP4-quantized gate weights."""
+    from dsv4.layers.linear import Nvfp4Linear
+    from dsv4.ops.quantize import quantize_weight_to_nvfp4
+    from dsv4.ops.layouts import make_b_k_major, assemble_raw_scales_2d3d_3d_side
+
+    w_fp4, w_sf, gsb_scalar = quantize_weight_to_nvfp4(W_gate_bf16)
+
+    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
+    # Feed raw weights into Nvfp4Linear the same way single_shot_inference does
+    gate_lin.fp4 = [w_fp4]
+    gate_lin.sf = [w_sf]
+    gate_lin.gs = [gsb_scalar.item()]  # global scale = gsb (will be folded by finalize)
+    gate_lin.ws2 = [None]  # no weight_scale_2 for random weights
+    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
+    gate_lin.finalize_weights()
+    return gate_lin
+
+
 def test_nvfp4_fused_router():
     """Test the NVFP4 fused router kernel against reference path.
 
@@ -107,35 +120,22 @@ def test_nvfp4_fused_router():
     """
     torch.manual_seed(42)
     device = "cuda"
-    M = 1        # decode: single token
-    K = 7168     # hidden size
-    N = 384      # num experts
+    M = 1
+    K = 7168
+    N = 384
     top_k = 6
     routed_scaling_factor = 0.5
 
     print(f"NVFP4 fused router test (M={M}, K={K}, N={N}, top_k={top_k}):")
 
-    # Create BF16 hidden states
     hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
-    # Create BF16 gate weight and quantize to NVFP4
     W_gate_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device)
     e_bias = torch.randn(N, dtype=torch.float32, device=device)
 
-    # Quantize weight to NVFP4 (matching Nvfp4Linear's quantization)
-    from dsv4.ops.quantize import quantize_weight_nvfp4, quantize_activation_nvfp4
-    w_fp4, w_sf, ws2_scalar = quantize_weight_nvfp4(W_gate_bf16)
+    # Build Nvfp4Linear with NVFP4 quantized weights
+    gate_lin = _build_nvfp4_linear(K, N, W_gate_bf16, device)
 
     # Reference path: NVFP4 GEMM + activation_topk
-    # Build Nvfp4Linear the same way single_shot_inference does
-    from dsv4.layers.linear import Nvfp4Linear
-    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
-    gate_lin.fp4 = [w_fp4]
-    gate_lin.sf = [w_sf]
-    gate_lin.gs = [1.0]
-    gate_lin.ws2 = [ws2_scalar.to(device) if ws2_scalar is not None else None]
-    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
-    gate_lin.finalize_weights()
-
     logits_ref = gate_lin(hidden_states).float()
     from dsv4.kernels.router._activation_topk import run_fused_activation_topk
     ref_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
@@ -144,7 +144,7 @@ def test_nvfp4_fused_router():
 
     print(f"  Reference: IDs={ref_ids[0].tolist()}, weights={[f'{w:.4f}' for w in ref_w[0].tolist()]}")
 
-    # Fused kernel path — use Nvfp4Linear's processed weight tensors
+    # Fused kernel path
     from dsv4.kernels.router.nvfp4_fused_router_kernel import run_nvfp4_fused_router
     gsb_val = gate_lin._gsb.item()
     gsa = gate_lin._activation_global_scale
@@ -155,7 +155,6 @@ def test_nvfp4_fused_router():
 
     print(f"  Fused:     IDs={fused_ids[0].tolist()}, weights={[f'{w:.4f}' for w in fused_w[0].tolist()]}")
 
-    # Compare IDs
     ids_match = (fused_ids == ref_ids).all().item()
     if ids_match:
         print(f"  IDs match reference: OK")
@@ -163,7 +162,6 @@ def test_nvfp4_fused_router():
         mismatches = (fused_ids != ref_ids).sum().item()
         print(f"  WARNING: {mismatches} ID mismatches vs reference")
 
-    # Compare weights
     if fused_w.shape == ref_w.shape:
         cos = torch.nn.functional.cosine_similarity(
             fused_w.flatten().unsqueeze(0),
@@ -178,7 +176,6 @@ def test_nvfp4_fused_router():
         else:
             print(f"  Weight match: POOR - needs investigation")
 
-    # Verify weight normalization
     w_sum = fused_w.sum(dim=1)
     for row in range(M):
         diff = abs(w_sum[row].item() - routed_scaling_factor)
@@ -204,18 +201,9 @@ def test_nvfp4_fused_router_multitoken():
     W_gate_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device)
     e_bias = torch.randn(N, dtype=torch.float32, device=device)
 
-    from dsv4.ops.quantize import quantize_weight_nvfp4
-    w_fp4, w_sf, ws2_scalar = quantize_weight_nvfp4(W_gate_bf16)
+    gate_lin = _build_nvfp4_linear(K, N, W_gate_bf16, device)
 
     # Reference
-    from dsv4.layers.linear import Nvfp4Linear
-    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
-    gate_lin.fp4 = [w_fp4]
-    gate_lin.sf = [w_sf]
-    gate_lin.gs = [1.0]
-    gate_lin.ws2 = [ws2_scalar.to(device) if ws2_scalar is not None else None]
-    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
-    gate_lin.finalize_weights()
     logits_ref = gate_lin(hidden_states).float()
     from dsv4.kernels.router._activation_topk import run_fused_activation_topk
     ref_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
@@ -231,7 +219,6 @@ def test_nvfp4_fused_router_multitoken():
         gsa, gsb_val, e_bias, routed_scaling_factor, top_k,
     )
 
-    # Compare
     ids_match = (fused_ids == ref_ids).all().item()
     mismatches = (fused_ids != ref_ids).sum().item() if not ids_match else 0
     cos = torch.nn.functional.cosine_similarity(
@@ -245,13 +232,16 @@ def test_nvfp4_fused_router_multitoken():
 if __name__ == "__main__":
     test_reference_router()
     print()
-    # NVFP4 fused tests require B200 — they'll fail on non-SM100 hardware
     try:
         test_nvfp4_fused_router()
     except Exception as e:
-        print(f"NVFP4 fused router test skipped (requires B200): {e}")
+        import traceback
+        traceback.print_exc()
+        print(f"NVFP4 fused router test failed: {e}")
     print()
     try:
         test_nvfp4_fused_router_multitoken()
     except Exception as e:
-        print(f"NVFP4 fused router multi-token test skipped (requires B200): {e}")
+        import traceback
+        traceback.print_exc()
+        print(f"NVFP4 fused router multi-token test failed: {e}")