[AMD] Fuse topk padded-token masking into a single Triton kernel (#28084)

Co-authored-by: Cursor <cursoragent@cursor.com>

python/sglang/srt/layers/moe/topk.py

@@ -30,6 +30,8 @@ from typing import (
 
 import torch
 import torch.nn.functional as F
+import triton
+import triton.language as tl
 
 from sglang.srt.runtime_context import get_parallel
 
@@ -130,6 +132,13 @@ _is_xpu = is_xpu()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 _is_musa = is_musa()
 
+# Experimental: skip the HIP padded-token routing-weight masking entirely.
+# Padded (CUDA-graph) rows are discarded downstream and the MoE combine is
+# per-token, so zeroing their weights is in principle unnecessary. Gated off by
+# default because it is a numerics-affecting change that must be validated with
+# an accuracy run before becoming the default.
+_skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False")
+
 if _is_cuda:
     from sgl_kernel import moe_fused_gate
 
@@ -1134,6 +1143,71 @@ def is_power_of_two(n):
     return n > 0 and math.log2(n).is_integer()
 
 
+@triton.jit
+def _fill_padded_rows_kernel(
+    out_ptr,
+    num_token_non_padded_ptr,
+    n_cols,
+    fill_value,
+    stride_row,
+    BLOCK_COLS: tl.constexpr,
+):
+    row = tl.program_id(0)
+    n_valid = tl.load(num_token_non_padded_ptr)
+    if row >= n_valid:
+        cols = tl.arange(0, BLOCK_COLS)
+        mask = cols < n_cols
+        ptrs = out_ptr + row * stride_row + cols
+        fill = tl.full((BLOCK_COLS,), fill_value, dtype=out_ptr.dtype.element_ty)
+        tl.store(ptrs, fill, mask=mask)
+
+
+def _can_fuse_padded_region(x: torch.Tensor) -> bool:
+    # The fused kernel uses one program per row and assumes a row-major 2D
+    # tensor (columns contiguous); fall back to eager for anything else.
+    return x.dim() == 2 and x.stride(1) == 1
+
+
+def _fill_padded_rows(
+    x: torch.Tensor,
+    num_token_non_padded: torch.Tensor,
+    fill_value,
+) -> None:
+    """Set ``x[row, :] = fill_value`` for every padded row (row index
+    ``>= num_token_non_padded``) using a single Triton launch.
+
+    Replaces the eager ``arange + (>=) + boolean index_put_`` sequence, which
+    issues several launch-latency-bound kernels per call. The grid is static
+    (one program per row) and the pad count is read from device memory inside
+    the kernel, so this is safe to capture inside a CUDA/HIP graph.
+    """
+    # Metadata-only checks (no device sync): the kernel reads a single scalar
+    # routing count from device memory, so it must be a 1-element integer tensor
+    # on the same device as ``x``.
+    assert isinstance(
+        num_token_non_padded, torch.Tensor
+    ), "num_token_non_padded must be a torch.Tensor"
+    assert num_token_non_padded.numel() == 1, (
+        "num_token_non_padded must be a single-element tensor, got shape "
+        f"{tuple(num_token_non_padded.shape)}"
+    )
+    assert (
+        not num_token_non_padded.dtype.is_floating_point
+    ), f"num_token_non_padded must be an integer tensor, got {num_token_non_padded.dtype}"
+    assert (
+        num_token_non_padded.device == x.device
+    ), "num_token_non_padded and x must be on the same device"
+    n_rows, n_cols = x.shape
+    _fill_padded_rows_kernel[(n_rows,)](
+        x,
+        num_token_non_padded,
+        n_cols,
+        fill_value,
+        x.stride(0),
+        BLOCK_COLS=triton.next_power_of_2(n_cols),
+    )
+
+
 def _eplb_remap_enabled() -> bool:
     # A real logical->physical mapping only exists when EPLB is enabled, the
     # initial expert placement is non-trivial, or there are redundant physical
@@ -1172,6 +1246,8 @@ def _mask_topk_ids_padded_region(
         indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
         mask = (indices >= num_token_non_padded).unsqueeze(-1)
         topk_ids = torch.where(mask, torch.full_like(topk_ids, -1), topk_ids)
+    elif _can_fuse_padded_region(topk_ids):
+        _fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
     else:
         indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
         topk_ids[indices >= num_token_non_padded, :] = fill_value
@@ -1183,6 +1259,9 @@ def _zero_topk_weights_padded_region(
 ):
     if num_token_non_padded is None:
         return
+    if _can_fuse_padded_region(topk_weights):
+        _fill_padded_rows(topk_weights, num_token_non_padded, 0.0)
+        return
     indices = torch.arange(0, topk_weights.shape[0], device=topk_weights.device)
     topk_weights[indices >= num_token_non_padded, :] = 0.0
 
@@ -1591,9 +1670,10 @@ def _post_process_topk_ids(
             topk_ids = topk_ids_logical_to_physical(
                 topk_ids, expert_location_dispatch_info
             )
-        # On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely, so
+        # NOTE (HIP): padded-token routing-weight zeroing is deferred to the
-        # padded tokens are neutralized by zeroing their routing weights.
+        # single pass at the end of this function (gated by SGLANG_MORI_NO_PAD_MASK).
-        _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
+        # That final pass re-zeros after any shared-expert append/remap, so a
+        # second zeroing here would be redundant (zeroing is idempotent).
 
     if recorder_topk_ids is None:
         recorder_topk_ids = topk_ids
@@ -1635,7 +1715,7 @@ def _post_process_topk_ids(
             topk_config,
         )
 
-    if _is_hip:
+    if _is_hip and not _skip_hip_pad_mask:
         # Shared-expert append/remap can introduce non-zero weights after the
         # initial HIP padding mask above. Ensure padded tokens leave this helper
         # with all expert weights zeroed.

test/registered/moe/test_topk_padded_region.py

@@ -0,0 +1,194 @@
+import unittest
+
+import torch
+
+import sglang.srt.layers.moe.topk as topk_mod
+from sglang.srt.layers.moe.topk import (
+    TopKConfig,
+    _can_fuse_padded_region,
+    _fill_padded_rows,
+    _mask_topk_ids_padded_region,
+    _post_process_topk_ids,
+    _zero_topk_weights_padded_region,
+)
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
+register_amd_ci(est_time=60, stage="stage-b", runner_config="1-gpu-small-amd")
+
+_IS_HIP = is_hip()
+
+torch.manual_seed(1234)
+
+
+def _eager_fill_padded_rows(x, num_token_non_padded, fill_value):
+    out = x.clone()
+    indices = torch.arange(0, x.shape[0], device=x.device)
+    out[indices >= num_token_non_padded, :] = fill_value
+    return out
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available(), "fused padded-region kernel needs a GPU"
+)
+class TestTopkPaddedRegion(CustomTestCase):
+    DEVICE = "cuda"
+
+    def test_matches_eager_across_shapes(self):
+        configs = [
+            # (n_tokens, topk, dtype, fill, helper)
+            (37, 9, torch.float32, 0.0, _zero_topk_weights_padded_region),
+            (256, 8, torch.float32, 0.0, _zero_topk_weights_padded_region),
+            (1, 9, torch.float32, 0.0, _zero_topk_weights_padded_region),
+            (37, 9, torch.int32, -1, _mask_topk_ids_padded_region),
+            (512, 16, torch.int32, -1, _mask_topk_ids_padded_region),
+        ]
+        for n, k, dtype, fill, helper in configs:
+            for n_valid in (0, 1, 5, n - 1, n):
+                with self.subTest(n=n, k=k, dtype=dtype, n_valid=n_valid):
+                    if dtype.is_floating_point:
+                        x = torch.rand((n, k), device=self.DEVICE, dtype=dtype) + 0.5
+                    else:
+                        x = torch.randint(
+                            0, 100, (n, k), device=self.DEVICE, dtype=dtype
+                        )
+                    self.assertTrue(_can_fuse_padded_region(x))
+                    num_token_non_padded = torch.tensor(
+                        n_valid, device=self.DEVICE, dtype=torch.int32
+                    )
+                    expected = _eager_fill_padded_rows(x, num_token_non_padded, fill)
+                    fused = x.clone()
+                    helper(fused, num_token_non_padded)
+                    self.assertTrue(torch.equal(fused, expected))
+
+    def test_none_pad_count_is_noop(self):
+        x = torch.rand((16, 8), device=self.DEVICE, dtype=torch.float32) + 0.5
+        ref = x.clone()
+        _zero_topk_weights_padded_region(x, None)
+        self.assertTrue(torch.equal(x, ref))
+
+    def test_non_contiguous_falls_back_to_eager(self):
+        # A column slice is not row-major contiguous, so the fused path must be
+        # skipped while still producing the correct result via the eager branch.
+        base = torch.rand((32, 16), device=self.DEVICE, dtype=torch.float32) + 0.5
+        view = base[:, ::2]
+        self.assertFalse(_can_fuse_padded_region(view))
+        num_token_non_padded = torch.tensor(5, device=self.DEVICE, dtype=torch.int32)
+        expected = _eager_fill_padded_rows(view, num_token_non_padded, 0.0)
+        _zero_topk_weights_padded_region(view, num_token_non_padded)
+        self.assertTrue(torch.equal(view, expected))
+
+    def test_cuda_graph_capture_and_replay(self):
+        n, k = 256, 9
+        weights = torch.rand((n, k), device=self.DEVICE, dtype=torch.float32) + 0.5
+        num_token_non_padded = torch.tensor(n, device=self.DEVICE, dtype=torch.int32)
+
+        # Warmup on a side stream before capture.
+        side = torch.cuda.Stream()
+        side.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(side):
+            for _ in range(3):
+                tmp = weights.clone()
+                _zero_topk_weights_padded_region(tmp, num_token_non_padded)
+        torch.cuda.current_stream().wait_stream(side)
+
+        work = weights.clone()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            _zero_topk_weights_padded_region(work, num_token_non_padded)
+
+        for n_valid in (n, 5, 0, 100):
+            work.copy_(weights)
+            num_token_non_padded.fill_(n_valid)
+            graph.replay()
+            torch.cuda.synchronize()
+            expected = _eager_fill_padded_rows(weights, num_token_non_padded, 0.0)
+            self.assertTrue(torch.equal(work, expected))
+
+    def test_invalid_pad_count_tensor_raises(self):
+        x = torch.rand((8, 8), device=self.DEVICE, dtype=torch.float32)
+        with self.assertRaises(AssertionError):
+            _fill_padded_rows(x, 4, 0.0)  # python int, not a tensor
+        with self.assertRaises(AssertionError):
+            _fill_padded_rows(
+                x,
+                torch.tensor([1, 2], device=self.DEVICE, dtype=torch.int32),
+                0.0,
+            )
+        with self.assertRaises(AssertionError):
+            _fill_padded_rows(
+                x,
+                torch.tensor(4.0, device=self.DEVICE, dtype=torch.float32),
+                0.0,
+            )
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "padded-region masking needs a GPU")
+class TestZeroPaddedRegionIdempotent(CustomTestCase):
+    """The HIP post-process keeps a single padded-row zeroing pass. Removing the
+    earlier (redundant) pass is only safe if zeroing is idempotent."""
+
+    DEVICE = "cuda"
+
+    def test_zeroing_twice_equals_once(self):
+        for n, k in [(37, 9), (256, 8), (512, 16)]:
+            for n_valid in (0, 5, n - 1, n):
+                with self.subTest(n=n, k=k, n_valid=n_valid):
+                    base = torch.rand((n, k), device=self.DEVICE) + 0.5
+                    pad = torch.tensor(n_valid, device=self.DEVICE, dtype=torch.int32)
+
+                    once = base.clone()
+                    _zero_topk_weights_padded_region(once, pad)
+
+                    twice = base.clone()
+                    _zero_topk_weights_padded_region(twice, pad)
+                    _zero_topk_weights_padded_region(twice, pad)
+
+                    self.assertTrue(torch.equal(once, twice))
+
+
+@unittest.skipUnless(
+    _IS_HIP and torch.cuda.is_available(),
+    "_post_process_topk_ids padded masking is HIP-only",
+)
+class TestPostProcessPaddedMaskingHip(CustomTestCase):
+    DEVICE = "cuda"
+
+    def _run(self, n=256, k=8, n_valid=5):
+        topk_weights = torch.rand((n, k), device=self.DEVICE, dtype=torch.float32) + 0.5
+        topk_ids = torch.randint(0, 64, (n, k), device=self.DEVICE, dtype=torch.int32)
+        router_logits = torch.rand((n, 64), device=self.DEVICE, dtype=torch.float32)
+        pad = torch.tensor(n_valid, device=self.DEVICE, dtype=torch.int32)
+        cfg = TopKConfig(top_k=k, num_fused_shared_experts=0)
+        _, out_weights, _ = _post_process_topk_ids(
+            topk_ids,
+            topk_weights,
+            cfg,
+            router_logits,
+            layer_id=0,
+            num_token_non_padded=pad,
+        )
+        return out_weights, n_valid
+
+    def test_padded_rows_zeroed_by_default(self):
+        # Flag off (default): padded rows must be fully zeroed, valid rows kept.
+        self.assertFalse(topk_mod._skip_hip_pad_mask)
+        out, n_valid = self._run()
+        self.assertTrue(torch.all(out[n_valid:] == 0.0))
+        self.assertTrue(torch.all(out[:n_valid] > 0.0))
+
+    def test_flag_skips_masking(self):
+        # Flag on: padded rows are left untouched (kept non-zero here).
+        orig = topk_mod._skip_hip_pad_mask
+        topk_mod._skip_hip_pad_mask = True
+        try:
+            out, n_valid = self._run()
+            self.assertTrue(torch.all(out[n_valid:] > 0.0))
+        finally:
+            topk_mod._skip_hip_pad_mask = orig
+
+
+if __name__ == "__main__":
+    unittest.main()