[MoE] Make simulated expert routing support DP>1, and fuse into one triton kernel (#29718)

Co-authored-by: jonnykong <jonnykong@fb.com>

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

@@ -9,6 +9,7 @@ from sglang.srt.layers.moe.utils import (
     get_moe_runner_backend,
     get_tbo_token_distribution_threshold,
     initialize_moe_config,
+    is_moe_input_scattered_across_dp_ranks,
     is_tbo_enabled,
     should_skip_mlp_all_reduce,
     should_skip_post_experts_all_reduce,
@@ -30,6 +31,7 @@ __all__ = [
     "should_skip_post_experts_all_reduce",
     "should_use_dp_reduce_scatterv",
     "should_use_flashinfer_cutlass_moe_fp4_allgather",
+    "is_moe_input_scattered_across_dp_ranks",
     "is_tbo_enabled",
     "get_tbo_token_distribution_threshold",
     "get_deepep_config",

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

@@ -31,6 +31,8 @@ from typing import (
 
 import torch
 import torch.nn.functional as F
+import triton
+import triton.language as tl
 
 if TYPE_CHECKING:
     from triton_kernels.tensor_details.ragged_tensor import RaggedTensorMetadata
@@ -93,10 +95,11 @@ from sglang.srt.eplb.expert_location_dispatch import (
     topk_ids_logical_to_physical,
 )
 from sglang.srt.layers.dp_attention import is_allocation_symmetric
-from sglang.srt.layers.moe import get_moe_runner_backend
+from sglang.srt.layers.moe import (
-from sglang.srt.layers.moe.utils import (
+    get_moe_runner_backend,
-    has_per_rank_fused_shared_slots,
+    is_moe_input_scattered_across_dp_ranks,
 )
+from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots
 from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -367,23 +370,116 @@ class PackedTopKOutput(NamedTuple):
         return TopKOutputFormat.PACKED
 
 
-def _make_round_robin_expert_ids(
+@triton.jit
-    num_tokens: int,
+def _simulate_balanced_routing_kernel(
-    topk: int,
+    topk_ids_ptr,
+    topk_weights_ptr,
+    num_experts,
+    step,
+    inv_k,
+    seed,
+    layer_offset,
+    token_shard_rank,
+    num_token_shards,
+    stride_im,
+    stride_ik,
+    stride_wm,
+    stride_wk,
+    K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    RANDOM: tl.constexpr,
+):
+    """One program per token: overwrite its top-k row with a balanced expert
+    assignment and uniform ``1/k`` weights, in a single launch — so the
+    benchmark override barely perturbs routing/MoE timing vs. the non-simulated
+    path (instead of the ~5-7 small elementwise ops it replaces).
+
+    Shapes:
+    - ``topk_ids_ptr``: ``[num_tokens, K]`` (row-major; strides passed in),
+      overwritten in place
+    - ``topk_weights_ptr``: ``[num_tokens, K]`` (row-major; strides passed in),
+      overwritten in place
+
+    ``RANDOM=False`` is the deterministic round-robin base ``token + layer_offset``;
+    ``RANDOM=True`` is a random per-token base (uniform, balanced in expectation;
+    ``seed`` is a kernel arg, so it is baked at CUDA-graph capture and replays stay
+    balanced). Both spread the k experts by ``step`` and emit global expert ids
+    (any EP logical->physical remap happens later in ``_post_process_topk_ids``).
+    ``token_shard_rank`` and ``num_token_shards`` ensure scattered DP ranks generate
+    different expert assignments for their local tokens when DP > 1."""
+    t = tl.program_id(0)
+    global_t = t * num_token_shards + token_shard_rank
+    j = tl.arange(0, BLOCK_K)
+    mask = j < K
+    if RANDOM:
+        base = (tl.rand(seed, global_t) * num_experts).to(tl.int32)
+    else:
+        base = global_t + layer_offset
+    gid = (base + j * step) % num_experts
+    tl.store(topk_ids_ptr + t * stride_im + j * stride_ik, gid, mask=mask)
+    tl.store(
+        topk_weights_ptr + t * stride_wm + j * stride_wk,
+        tl.full((BLOCK_K,), inv_k, tl.float32),
+        mask=mask,
+    )
+
+
+# Per-launch seed for the uniform (RANDOM=True) path: varies across eager calls so
+# the random base differs, while being baked at CUDA-graph capture (graph-safe).
+_simulate_uniform_seed = 0
+
+
+def _simulate_balanced_routing(
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
     num_experts: int,
     *,
-    device: torch.device,
+    random: bool,
-    dtype: torch.dtype,
     layer_id: Optional[int] = None,
-) -> torch.Tensor:
+    token_shard_rank: int = 0,
-    if topk == 0:
+    num_token_shards: int = 1,
-        return torch.empty((num_tokens, 0), device=device, dtype=dtype)
+    seed: Optional[int] = None,
+) -> None:
+    """Benchmark-only fused override (in place): replace ``topk_ids`` with a
+    balanced expert assignment and ``topk_weights`` with ``1/k`` using a single
+    Triton kernel. ``random=False`` is round-robin; ``random=True`` is uniform.
 
-    step = max(num_experts // topk, 1)
+    Shapes:
-    layer_offset = 0 if layer_id is None else layer_id
+    - ``topk_ids``: ``[num_tokens, k]``, overwritten in place
-    offsets = torch.arange(num_tokens, device=device, dtype=dtype).unsqueeze(1)
+    - ``topk_weights``: ``[num_tokens, k]``, overwritten in place
-    steps = torch.arange(topk, device=device, dtype=dtype).unsqueeze(0) * step
+
-    return (offsets + layer_offset + steps) % num_experts
+    ``token_shard_rank`` and ``num_token_shards`` describe scattered DP input.
+    Their defaults describe a gathered token buffer (effective DP=1). ``seed``
+    is exposed for deterministic tests; production calls use a per-launch seed.
+    """
+    global _simulate_uniform_seed
+    num_tokens, k = topk_ids.shape
+    if num_tokens == 0 or k == 0:
+        return
+    assert 0 <= token_shard_rank < num_token_shards
+    if random and seed is None:
+        seed = _simulate_uniform_seed
+        _simulate_uniform_seed += 1
+    elif seed is None:
+        seed = 0
+    _simulate_balanced_routing_kernel[(num_tokens,)](
+        topk_ids,
+        topk_weights,
+        num_experts,
+        max(num_experts // k, 1),
+        1.0 / k,
+        seed,
+        0 if layer_id is None else layer_id,
+        token_shard_rank,
+        num_token_shards,
+        topk_ids.stride(0),
+        topk_ids.stride(1),
+        topk_weights.stride(0),
+        topk_weights.stride(1),
+        K=k,
+        BLOCK_K=triton.next_power_of_2(k),
+        RANDOM=random,
+    )
 
 
 # -------------------------------- TopK ---------------------------------------
@@ -2367,34 +2463,29 @@ def select_experts(
             "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS are mutually exclusive"
         )
 
-    if simulate_uniform_experts:
+    if simulate_uniform_experts or simulate_round_robin_experts:
-        # Benchmark-only: override gating with random-offset uniform expert assignment
+        # Benchmark-only: override gating with a balanced expert assignment (so
-        # to avoid expert imbalance from dummy/random weights. Do NOT use in production.
+        # dummy/random benchmark tokens don't skew MoE load) via a single fused
-        num_tokens, k = topk_ids.shape
+        # Triton kernel — one launch instead of the ~5-7 small elementwise ops it
-        num_experts = router_logits.shape[1]
+        # replaces, to minimize timing perturbation. Do NOT use in production.
-        if k > 0:
+        if is_moe_input_scattered_across_dp_ranks():
-            offsets = torch.randint(
+            parallel = get_parallel()
-                0, num_experts, (num_tokens, 1), device=topk_ids.device
+            token_shard_rank = parallel.attn_dp_rank
-            )
+            num_token_shards = parallel.attn_dp_size
-            steps = torch.arange(k, device=topk_ids.device).unsqueeze(0)
+        else:
-            step = max(num_experts // k, 1)
+            # Gathered MoE presents one global token buffer to every rank, so
-            topk_ids = ((offsets + steps * step) % num_experts).to(topk_ids.dtype)
+            # its routing must remain identical across those replicas.
-            topk_weights = torch.ones_like(topk_weights) / k
+            token_shard_rank, num_token_shards = 0, 1
-    elif simulate_round_robin_experts:
+
-        # Benchmark-only: override gating with deterministic expert assignment
+        _simulate_balanced_routing(
-        # to avoid routing noise from dummy/random weights. Do NOT use in production.
+            topk_ids,
-        num_tokens, k = topk_ids.shape
+            topk_weights,
-        num_experts = router_logits.shape[1]
+            router_logits.shape[1],
-        topk_ids = _make_round_robin_expert_ids(
+            random=simulate_uniform_experts,
-            num_tokens,
-            k,
-            num_experts,
-            device=topk_ids.device,
-            dtype=topk_ids.dtype,
             layer_id=layer_id,
+            token_shard_rank=token_shard_rank,
+            num_token_shards=num_token_shards,
         )
-        if k > 0:
-            topk_weights = torch.full_like(topk_weights, 1.0 / k)
 
     topk_ids, topk_weights, recorder_topk_ids = _post_process_topk_ids(
         topk_ids=topk_ids,

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

@@ -595,6 +595,15 @@ def should_use_flashinfer_cutlass_moe_fp4_allgather():
     )
 
 
+def is_moe_input_scattered_across_dp_ranks() -> bool:
+    """Whether sparse MoE routing runs on a DP-local token shard."""
+    return (
+        not get_moe_a2a_backend().is_none()
+        or should_use_flashinfer_cutlass_moe_fp4_allgather()
+        or get_parallel().dwdp_size > 1
+    )
+
+
 def should_use_dp_reduce_scatterv():
     """
     Use reduce_scatterv in the standard dispatcher's combine() for DP attention

test/manual/layers/moe/test_simulate_balanced_routing.py

@@ -0,0 +1,165 @@
+"""Unit tests for the fused benchmark-only balanced-routing override in
+``sglang.srt.layers.moe.topk`` (``_simulate_balanced_routing`` /
+``_simulate_balanced_routing_kernel``).
+
+Verifies the single fused Triton kernel reproduces the
+``_make_round_robin_expert_ids`` reference exactly (incl. the per-layer offset),
+writes uniform ``1/k`` weights, and that the uniform path is structurally
+balanced. GPU-only (skips without CUDA).
+
+Run:
+    python -m pytest test/manual/layers/moe/test_simulate_balanced_routing.py -v
+"""
+
+import unittest
+from typing import Optional, Tuple
+
+import torch
+from parameterized import parameterized
+
+from sglang.srt.layers.moe.topk import _simulate_balanced_routing
+from sglang.test.test_utils import CustomTestCase
+
+E = 256  # num_experts
+K = 8  # top-k
+
+
+def _make_round_robin_expert_ids(
+    num_tokens: int,
+    topk: int,
+    num_experts: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    layer_id: Optional[int] = None,
+) -> torch.Tensor:
+    # Deterministic, perfectly balanced expert assignment: each token's top-k is
+    # spread by num_experts//topk. Returns global expert ids of shape
+    # [num_tokens, topk].
+    if topk == 0:
+        return torch.empty((num_tokens, 0), device=device, dtype=dtype)
+
+    step = max(num_experts // topk, 1)
+    layer_offset = 0 if layer_id is None else layer_id
+    offsets = torch.arange(num_tokens, device=device, dtype=dtype).unsqueeze(
+        1
+    )  # [num_tokens, 1]
+    steps = (
+        torch.arange(topk, device=device, dtype=dtype).unsqueeze(0) * step
+    )  # [1, topk]
+    return (offsets + layer_offset + steps) % num_experts  # [num_tokens, topk]
+
+
+def _alloc(
+    num_tokens: int, k: int, device: str = "cuda"
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Pre-filled with junk so the test fails if the kernel doesn't overwrite.
+    ids = torch.full((num_tokens, k), -7, dtype=torch.int32, device=device)
+    weights = torch.full((num_tokens, k), -7.0, dtype=torch.float32, device=device)
+    return ids, weights
+
+
+class TestSimulateBalancedRouting(CustomTestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+
+    # round-robin output must equal the reference exactly, for several layer
+    # offsets and both a power-of-2 and a non-power-of-2 top-k (BLOCK_K masking).
+    @parameterized.expand(
+        [
+            ("layer0_k8", 0, 8),
+            ("layer5_k8", 5, 8),
+            ("noneLayer_k8", None, 8),
+            ("layer3_k6", 3, 6),
+        ]
+    )
+    def test_round_robin_matches_reference(
+        self, _name: str, layer_id: Optional[int], k: int
+    ) -> None:
+        T = 512
+        ids, weights = _alloc(T, k)
+        _simulate_balanced_routing(ids, weights, E, random=False, layer_id=layer_id)
+        ref = _make_round_robin_expert_ids(
+            T, k, E, device="cuda", dtype=torch.int32, layer_id=layer_id
+        )
+        self.assertTrue(torch.equal(ids, ref))
+        torch.testing.assert_close(weights, torch.full_like(weights, 1.0 / k))
+
+    def test_round_robin_perfectly_balanced(self) -> None:
+        T = 512  # multiple of E -> exactly uniform per-expert load
+        ids, weights = _alloc(T, K)
+        _simulate_balanced_routing(ids, weights, E, random=False, layer_id=0)
+        counts = torch.bincount(ids.flatten().long(), minlength=E)
+        self.assertTrue(torch.all(counts == (T * K // E)))
+        for row in ids:
+            self.assertEqual(row.unique().numel(), K)
+
+    def test_uniform_structural(self) -> None:
+        # uniform: random per-token base, so assert only seed-independent props.
+        T = 4096
+        ids, weights = _alloc(T, K)
+        _simulate_balanced_routing(ids, weights, E, random=True, layer_id=0)
+        torch.testing.assert_close(weights, torch.full_like(weights, 1.0 / K))
+        self.assertGreaterEqual(int(ids.min()), 0)
+        self.assertLess(int(ids.max()), E)
+        # offset + j*step spreads the k experts out -> k distinct per row
+        for row in ids[:64]:
+            self.assertEqual(row.unique().numel(), K)
+
+    @parameterized.expand(
+        [
+            ("round_robin_dp2", False, 2),
+            ("round_robin_dp4", False, 4),
+            ("uniform_dp2", True, 2),
+            ("uniform_dp4", True, 4),
+        ]
+    )
+    def test_interleaved_dp_assignments_match_dp1(
+        self, _name: str, random: bool, dp_size: int
+    ) -> None:
+        # Interleaving the DP-local outputs must exactly reproduce the expert
+        # assignments for the equivalent DP=1 input. The fixed seed models
+        # independent processes entering the same uniform-routing call with
+        # the same initial seed.
+        T = 16
+        seed = 17
+        layer_id = 3
+        ids_by_rank = []
+        weights_by_rank = []
+        for dp_rank in range(dp_size):
+            ids, weights = _alloc(T, K)
+            _simulate_balanced_routing(
+                ids,
+                weights,
+                E,
+                random=random,
+                layer_id=layer_id,
+                token_shard_rank=dp_rank,
+                num_token_shards=dp_size,
+                seed=seed,
+            )
+            ids_by_rank.append(ids)
+            weights_by_rank.append(weights)
+
+        interleaved_ids = torch.stack(ids_by_rank, dim=1).reshape(T * dp_size, K)
+        interleaved_weights = torch.stack(weights_by_rank, dim=1).reshape(
+            T * dp_size, K
+        )
+
+        expected_ids, expected_weights = _alloc(T * dp_size, K)
+        _simulate_balanced_routing(
+            expected_ids,
+            expected_weights,
+            E,
+            random=random,
+            layer_id=layer_id,
+            seed=seed,
+        )
+
+        self.assertTrue(torch.equal(interleaved_ids, expected_ids))
+        torch.testing.assert_close(interleaved_weights, expected_weights)
+
+
+if __name__ == "__main__":
+    unittest.main()