[MoE] Deprecate act_and_mul_triton; fold filter_expert into JIT silu/gelu_and_mul (#23707)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/jit_kernel/tests/test_activation.py

@@ -75,5 +75,85 @@ def test_activation_out_param(
     torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
 
 
+FILTER_SHAPES = get_ci_test_range(
+    full_range=[(83, 1024), (256, 2048), (1024, 4096)],
+    ci_range=[(83, 1024)],
+)
+EXPERT_STEPS = [1, 16]
+
+
+@pytest.mark.parametrize("op_name", OPS)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", FILTER_SHAPES)
+@pytest.mark.parametrize("expert_step", EXPERT_STEPS)
+def test_activation_filter_expert(
+    op_name: str,
+    dtype: torch.dtype,
+    shape: tuple[int, int],
+    expert_step: int,
+) -> None:
+    """expert_ids[token // expert_step] == -1 must leave the output row untouched."""
+    num_tokens = shape[0]
+    x = torch.randn(shape, dtype=dtype, device="cuda")
+    # Pre-fill out with a sentinel so we can detect untouched rows.
+    sentinel = float("nan")
+    out = torch.full(
+        shape[:-1] + (shape[-1] // 2,),
+        sentinel,
+        dtype=dtype,
+        device="cuda",
+    )
+
+    num_groups = (num_tokens + expert_step - 1) // expert_step
+    expert_ids = torch.randint(
+        low=0, high=8, size=(num_groups,), dtype=torch.int32, device="cuda"
+    )
+    skip_mask = torch.rand(num_groups, device="cuda") < 0.4
+    expert_ids[skip_mask] = -1
+
+    result = run_activation(op_name, x, out, expert_ids, expert_step)
+    assert result is out
+
+    token_skip = skip_mask[torch.arange(num_tokens, device="cuda") // expert_step]
+    expected = _reference(op_name, x)
+    atol, rtol = _tolerances(dtype)
+
+    kept = ~token_skip
+    if kept.any():
+        torch.testing.assert_close(out[kept], expected[kept], atol=atol, rtol=rtol)
+    if token_skip.any():
+        assert torch.isnan(
+            out[token_skip]
+        ).all(), "filter_expert kernel touched rows whose expert_id is -1"
+
+
+@pytest.mark.parametrize("op_name", OPS)
+def test_activation_filter_expert_all_skipped(op_name: str) -> None:
+    """If every expert id is -1, the output must be left entirely untouched."""
+    shape = (32, 512)
+    x = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    out = torch.full(
+        shape[:-1] + (shape[-1] // 2,),
+        float("nan"),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    expert_ids = torch.full((shape[0],), -1, dtype=torch.int32, device="cuda")
+    run_activation(op_name, x, out, expert_ids, 1)
+    assert torch.isnan(out).all()
+
+
+@pytest.mark.parametrize("op_name", OPS)
+def test_activation_filter_expert_none_skipped(op_name: str) -> None:
+    """No -1 in expert_ids must yield bit-identical output to the unfiltered path."""
+    shape = (64, 512)
+    dtype = torch.bfloat16
+    x = torch.randn(shape, dtype=dtype, device="cuda")
+    expert_ids = torch.zeros((shape[0],), dtype=torch.int32, device="cuda")
+    out_filtered = run_activation(op_name, x, None, expert_ids, 1)
+    out_unfiltered = run_activation(op_name, x, None)
+    torch.testing.assert_close(out_filtered, out_unfiltered, atol=0.0, rtol=0.0)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))