test: add scores and hidden_clamp tests for fused_mul_grouped_poly_norm

- Rename test file to match fused_mul_grouped_poly_norm convention
- Add forward/backward tests for scores fusion
- Add forward/backward tests for hidden_clamp fusion (clamp values 10.0, 1.0, 0.5)
- Adjust weight/bias grad tolerance for atomicAdd accumulation order
- 192 tests, all passing

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

tests/{test_grouped_fused_mul_poly_norm.py → test_fused_mul_grouped_poly_norm.py}

@@ -3,11 +3,11 @@ import torch
 
 from grouped_poly_norm import (
     HAS_TRITON,
-    grouped_fused_mul_poly_norm_ref,
+    fused_mul_grouped_poly_norm_ref,
 )
 
 if HAS_TRITON:
-    from grouped_poly_norm import grouped_fused_mul_poly_norm
+    from grouped_poly_norm import fused_mul_grouped_poly_norm
 
 from .utils import assert_close
 
@@ -26,7 +26,7 @@ CUDA_DEVICES = ["cuda:0"]
 def _counts_to_offsets(counts_list, device):
     """Convert list of counts to cumsum offsets tensor."""
     return torch.cumsum(
-        torch.tensor(counts_list, device=device, dtype=torch.int32), dim=0)
+        torch.tensor(counts_list, device=device, dtype=torch.int32), dim=0).to(torch.int32)
 
 
 def _make_inputs(total_tokens, hidden_dim, num_experts, dtype, device,
@@ -54,32 +54,45 @@ def _make_inputs(total_tokens, hidden_dim, num_experts, dtype, device,
     return input_t, mul_t, weight, bias, offsets
 
 
-def _run_ref(input_t, mul_t, weight, bias, offsets, expert_offset=0):
+def _make_scores(total_tokens, device, dtype=torch.float32):
+    """Create random scores (N, 1) in fp32."""
+    return torch.rand(total_tokens, 1, device=device, dtype=dtype) * 0.5 + 0.5
+
+
+def _run_ref(input_t, mul_t, weight, bias, offsets, expert_offset=0,
+             scores=None, hidden_clamp=None):
     """Run reference forward + backward, return output and grads."""
     inp = input_t.clone().detach().requires_grad_(True)
     m = mul_t.clone().detach().requires_grad_(True)
     w = weight.clone().detach().requires_grad_(True)
     b = bias.clone().detach().requires_grad_(True)
+    s = scores.clone().detach().requires_grad_(True) if scores is not None else None
 
-    out = grouped_fused_mul_poly_norm_ref(inp, m, w, b, offsets,
-                                          expert_offset=expert_offset)
+    out = fused_mul_grouped_poly_norm_ref(inp, m, w, b, offsets,
+                                          expert_offset=expert_offset,
+                                          scores=s, hidden_clamp=hidden_clamp)
     out.sum().backward()
 
-    return out, inp.grad, m.grad, w.grad, b.grad
+    grads = (out, inp.grad, m.grad, w.grad, b.grad)
+    return grads + (s.grad,) if s is not None else grads + (None,)
 
 
-def _run_triton(input_t, mul_t, weight, bias, offsets, expert_offset=0):
-    """Run Triton forward + backward, return output and grads."""
+def _run_triton(input_t, mul_t, weight, bias, offsets, expert_offset=0,
+                scores=None, hidden_clamp=None):
+    """Run Triton/CUDA forward + backward, return output and grads."""
     inp = input_t.clone().detach().requires_grad_(True)
     m = mul_t.clone().detach().requires_grad_(True)
     w = weight.clone().detach().requires_grad_(True)
     b = bias.clone().detach().requires_grad_(True)
+    s = scores.clone().detach().requires_grad_(True) if scores is not None else None
 
-    out = grouped_fused_mul_poly_norm(inp, m, w, b, offsets,
-                                      expert_offset=expert_offset)
+    out = fused_mul_grouped_poly_norm(inp, m, w, b, offsets,
+                                      expert_offset=expert_offset,
+                                      scores=s, hidden_clamp=hidden_clamp)
     out.sum().backward()
 
-    return out, inp.grad, m.grad, w.grad, b.grad
+    grads = (out, inp.grad, m.grad, w.grad, b.grad)
+    return grads + (s.grad,) if s is not None else grads + (None,)
 
 
 @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
@@ -90,7 +103,7 @@ def _run_triton(input_t, mul_t, weight, bias, offsets, expert_offset=0):
 @pytest.mark.parametrize("expert_offset", EXPERT_OFFSETS)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-def test_grouped_fused_mul_poly_norm_forward(
+def test_fused_mul_grouped_poly_norm_forward(
     num_tokens: int,
     d: int,
     num_experts: int,
@@ -105,10 +118,10 @@ def test_grouped_fused_mul_poly_norm_forward(
         num_tokens, d, num_experts, dtype, device, seed,
         expert_offset=expert_offset)
 
-    out_ref = grouped_fused_mul_poly_norm_ref(input_t, mul_t, weight, bias,
+    out_ref = fused_mul_grouped_poly_norm_ref(input_t, mul_t, weight, bias,
                                               offsets,
                                               expert_offset=expert_offset)
-    out_tri = grouped_fused_mul_poly_norm(input_t, mul_t, weight, bias,
+    out_tri = fused_mul_grouped_poly_norm(input_t, mul_t, weight, bias,
                                          offsets,
                                          expert_offset=expert_offset)
 
@@ -129,7 +142,7 @@ def test_grouped_fused_mul_poly_norm_forward(
 @pytest.mark.parametrize("expert_offset", EXPERT_OFFSETS)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-def test_grouped_fused_mul_poly_norm_backward(
+def test_fused_mul_grouped_poly_norm_backward(
     num_tokens: int,
     d: int,
     num_experts: int,
@@ -144,9 +157,9 @@ def test_grouped_fused_mul_poly_norm_backward(
         num_tokens, d, num_experts, dtype, device, seed,
         expert_offset=expert_offset)
 
-    _, inp_grad_ref, mul_grad_ref, w_grad_ref, b_grad_ref = _run_ref(
+    _, inp_grad_ref, mul_grad_ref, w_grad_ref, b_grad_ref, _ = _run_ref(
         input_t, mul_t, weight, bias, offsets, expert_offset=expert_offset)
-    _, inp_grad_tri, mul_grad_tri, w_grad_tri, b_grad_tri = _run_triton(
+    _, inp_grad_tri, mul_grad_tri, w_grad_tri, b_grad_tri, _ = _run_triton(
         input_t, mul_t, weight, bias, offsets, expert_offset=expert_offset)
 
     if dtype == torch.float32:
@@ -164,7 +177,7 @@ def test_grouped_fused_mul_poly_norm_backward(
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("expert_offset", EXPERT_OFFSETS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-def test_grouped_fused_mul_poly_norm_zero_token_experts(
+def test_fused_mul_grouped_poly_norm_zero_token_experts(
     dtype: torch.dtype,
     expert_offset: int,
     device: str,
@@ -184,10 +197,10 @@ def test_grouped_fused_mul_poly_norm_zero_token_experts(
     bias = torch.zeros(total_experts, 1, device=device, dtype=dtype)
     offsets = _counts_to_offsets(counts, device)
 
-    out_ref = grouped_fused_mul_poly_norm_ref(input_t, mul_t, weight, bias,
+    out_ref = fused_mul_grouped_poly_norm_ref(input_t, mul_t, weight, bias,
                                               offsets,
                                               expert_offset=expert_offset)
-    out_tri = grouped_fused_mul_poly_norm(input_t, mul_t, weight, bias,
+    out_tri = fused_mul_grouped_poly_norm(input_t, mul_t, weight, bias,
                                          offsets,
                                          expert_offset=expert_offset)
 
@@ -197,12 +210,12 @@ def test_grouped_fused_mul_poly_norm_zero_token_experts(
         assert_close(out_ref, out_tri, atol=1e-2, rtol=1e-2)
 
     # Check backward with zero-token experts
-    _, _, _, w_grad_ref, b_grad_ref = _run_ref(input_t, mul_t, weight, bias,
-                                               offsets,
-                                               expert_offset=expert_offset)
-    _, _, _, w_grad_tri, b_grad_tri = _run_triton(input_t, mul_t, weight, bias,
-                                                  offsets,
-                                                  expert_offset=expert_offset)
+    _, _, _, w_grad_ref, b_grad_ref, _ = _run_ref(input_t, mul_t, weight, bias,
+                                                   offsets,
+                                                   expert_offset=expert_offset)
+    _, _, _, w_grad_tri, b_grad_tri, _ = _run_triton(input_t, mul_t, weight, bias,
+                                                      offsets,
+                                                      expert_offset=expert_offset)
 
     if dtype == torch.float32:
         atol, rtol = 1e-3, 1e-3
@@ -227,7 +240,7 @@ def test_grouped_fused_mul_poly_norm_zero_token_experts(
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("expert_offset", EXPERT_OFFSETS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-def test_grouped_fused_mul_poly_norm_no_nan_inf(
+def test_fused_mul_grouped_poly_norm_no_nan_inf(
     dtype: torch.dtype,
     expert_offset: int,
     device: str,
@@ -237,7 +250,7 @@ def test_grouped_fused_mul_poly_norm_no_nan_inf(
     input_t, mul_t, weight, bias, offsets = _make_inputs(
         4096, 256, 8, dtype, device, expert_offset=expert_offset)
 
-    out, inp_grad, mul_grad, w_grad, b_grad = _run_triton(
+    out, inp_grad, mul_grad, w_grad, b_grad, _ = _run_triton(
         input_t, mul_t, weight, bias, offsets, expert_offset=expert_offset)
 
     assert not out.isnan().any(), "Output contains NaN"
@@ -247,3 +260,128 @@ def test_grouped_fused_mul_poly_norm_no_nan_inf(
                        ("weight", w_grad), ("bias", b_grad)]:
         assert not grad.isnan().any(), f"{name}_grad contains NaN"
         assert not grad.isinf().any(), f"{name}_grad contains Inf"
+
+
+# ---------------------------------------------------------------------------
+# Scores tests
+# ---------------------------------------------------------------------------
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("num_experts", [8, 48])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_fused_mul_grouped_poly_norm_scores_forward(
+    num_tokens, d, num_experts, dtype, device,
+):
+    """Forward with scores should match reference."""
+    torch.set_default_device(device)
+    input_t, mul_t, weight, bias, offsets = _make_inputs(
+        num_tokens, d, num_experts, dtype, device)
+    scores = _make_scores(num_tokens, device)
+
+    out_ref = fused_mul_grouped_poly_norm_ref(
+        input_t, mul_t, weight, bias, offsets, scores=scores)
+    out_tri = fused_mul_grouped_poly_norm(
+        input_t, mul_t, weight, bias, offsets, scores=scores)
+
+    atol, rtol = (1e-4, 1e-4) if dtype == torch.float32 else (1e-2, 1e-2)
+    assert_close(out_ref, out_tri, atol=atol, rtol=rtol)
+
+
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("num_experts", [8, 48])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_fused_mul_grouped_poly_norm_scores_backward(
+    num_tokens, d, num_experts, dtype, device,
+):
+    """Backward with scores should match reference."""
+    torch.set_default_device(device)
+    input_t, mul_t, weight, bias, offsets = _make_inputs(
+        num_tokens, d, num_experts, dtype, device)
+    scores = _make_scores(num_tokens, device)
+
+    out_ref, ig_ref, mg_ref, wg_ref, bg_ref, sg_ref = _run_ref(
+        input_t, mul_t, weight, bias, offsets, scores=scores)
+    out_tri, ig_tri, mg_tri, wg_tri, bg_tri, sg_tri = _run_triton(
+        input_t, mul_t, weight, bias, offsets, scores=scores)
+
+    atol, rtol = (1e-4, 1e-4) if dtype == torch.float32 else (5e-2, 5e-2)
+    # weight/bias grads use atomicAdd accumulation across tokens,
+    # so allow slightly higher tolerance for fp32
+    wg_atol = 5e-4 if dtype == torch.float32 else 5e-2
+    assert_close(ig_ref, ig_tri, atol=atol, rtol=rtol)
+    assert_close(mg_ref, mg_tri, atol=atol, rtol=rtol)
+    assert_close(wg_ref, wg_tri, atol=wg_atol, rtol=wg_atol)
+    assert_close(bg_ref, bg_tri, atol=wg_atol, rtol=wg_atol)
+    assert_close(sg_ref, sg_tri, atol=atol, rtol=rtol)
+
+
+# ---------------------------------------------------------------------------
+# Hidden clamp tests
+# ---------------------------------------------------------------------------
+CLAMP_VALUES = [10.0, 1.0, 0.5]
+
+
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
+@pytest.mark.parametrize("num_tokens", [4096])
+@pytest.mark.parametrize("d", [256, 1280])
+@pytest.mark.parametrize("num_experts", [8])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("hidden_clamp", CLAMP_VALUES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_fused_mul_grouped_poly_norm_hidden_clamp_forward(
+    num_tokens, d, num_experts, dtype, hidden_clamp, device,
+):
+    """Forward with hidden_clamp should match reference."""
+    torch.set_default_device(device)
+    input_t, mul_t, weight, bias, offsets = _make_inputs(
+        num_tokens, d, num_experts, dtype, device)
+    scores = _make_scores(num_tokens, device)
+
+    out_ref = fused_mul_grouped_poly_norm_ref(
+        input_t, mul_t, weight, bias, offsets,
+        scores=scores, hidden_clamp=hidden_clamp)
+    out_tri = fused_mul_grouped_poly_norm(
+        input_t, mul_t, weight, bias, offsets,
+        scores=scores, hidden_clamp=hidden_clamp)
+
+    atol, rtol = (1e-4, 1e-4) if dtype == torch.float32 else (1e-2, 1e-2)
+    assert_close(out_ref, out_tri, atol=atol, rtol=rtol)
+
+
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
+@pytest.mark.parametrize("num_tokens", [4096])
+@pytest.mark.parametrize("d", [256, 1280])
+@pytest.mark.parametrize("num_experts", [8])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("hidden_clamp", CLAMP_VALUES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_fused_mul_grouped_poly_norm_hidden_clamp_backward(
+    num_tokens, d, num_experts, dtype, hidden_clamp, device,
+):
+    """Backward with hidden_clamp should match reference."""
+    torch.set_default_device(device)
+    input_t, mul_t, weight, bias, offsets = _make_inputs(
+        num_tokens, d, num_experts, dtype, device)
+    scores = _make_scores(num_tokens, device)
+
+    out_ref, ig_ref, mg_ref, wg_ref, bg_ref, sg_ref = _run_ref(
+        input_t, mul_t, weight, bias, offsets,
+        scores=scores, hidden_clamp=hidden_clamp)
+    out_tri, ig_tri, mg_tri, wg_tri, bg_tri, sg_tri = _run_triton(
+        input_t, mul_t, weight, bias, offsets,
+        scores=scores, hidden_clamp=hidden_clamp)
+
+    atol, rtol = (1e-4, 1e-4) if dtype == torch.float32 else (5e-2, 5e-2)
+    # weight/bias grads use atomicAdd accumulation across tokens,
+    # so allow slightly higher tolerance for fp32
+    wg_atol = 5e-4 if dtype == torch.float32 else 5e-2
+    assert_close(ig_ref, ig_tri, atol=atol, rtol=rtol)
+    assert_close(mg_ref, mg_tri, atol=atol, rtol=rtol)
+    assert_close(wg_ref, wg_tri, atol=wg_atol, rtol=wg_atol)
+    assert_close(bg_ref, bg_tri, atol=wg_atol, rtol=wg_atol)
+    assert_close(sg_ref, sg_tri, atol=atol, rtol=rtol)