[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.utils import (
-    has_per_rank_fused_shared_slots,
+from sglang.srt.layers.moe import (
+    get_moe_runner_backend,
+    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(
-    num_tokens: int,
-    topk: int,
+@triton.jit
+def _simulate_balanced_routing_kernel(
+    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,
-    dtype: torch.dtype,
+    random: bool,
     layer_id: Optional[int] = None,
-) -> torch.Tensor:
-    if topk == 0:
-        return torch.empty((num_tokens, 0), device=device, dtype=dtype)
+    token_shard_rank: int = 0,
+    num_token_shards: int = 1,
+    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)
-    layer_offset = 0 if layer_id is None else layer_id
-    offsets = torch.arange(num_tokens, device=device, dtype=dtype).unsqueeze(1)
-    steps = torch.arange(topk, device=device, dtype=dtype).unsqueeze(0) * step
-    return (offsets + layer_offset + steps) % num_experts
+    Shapes:
+    - ``topk_ids``: ``[num_tokens, k]``, overwritten in place
+    - ``topk_weights``: ``[num_tokens, k]``, overwritten in place
+
+    ``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:
-        # Benchmark-only: override gating with random-offset uniform expert assignment
-        # to avoid expert imbalance from dummy/random weights. Do NOT use in production.
-        num_tokens, k = topk_ids.shape
-        num_experts = router_logits.shape[1]
-        if k > 0:
-            offsets = torch.randint(
-                0, num_experts, (num_tokens, 1), device=topk_ids.device
-            )
-            steps = torch.arange(k, device=topk_ids.device).unsqueeze(0)
-            step = max(num_experts // k, 1)
-            topk_ids = ((offsets + steps * step) % num_experts).to(topk_ids.dtype)
-            topk_weights = torch.ones_like(topk_weights) / k
-    elif simulate_round_robin_experts:
-        # Benchmark-only: override gating with deterministic expert assignment
-        # to avoid routing noise from dummy/random weights. Do NOT use in production.
-        num_tokens, k = topk_ids.shape
-        num_experts = router_logits.shape[1]
-        topk_ids = _make_round_robin_expert_ids(
-            num_tokens,
-            k,
-            num_experts,
-            device=topk_ids.device,
-            dtype=topk_ids.dtype,
+    if simulate_uniform_experts or simulate_round_robin_experts:
+        # Benchmark-only: override gating with a balanced expert assignment (so
+        # dummy/random benchmark tokens don't skew MoE load) via a single fused
+        # Triton kernel — one launch instead of the ~5-7 small elementwise ops it
+        # replaces, to minimize timing perturbation. Do NOT use in production.
+        if is_moe_input_scattered_across_dp_ranks():
+            parallel = get_parallel()
+            token_shard_rank = parallel.attn_dp_rank
+            num_token_shards = parallel.attn_dp_size
+        else:
+            # Gathered MoE presents one global token buffer to every rank, so
+            # its routing must remain identical across those replicas.
+            token_shard_rank, num_token_shards = 0, 1
+
+        _simulate_balanced_routing(
+            topk_ids,
+            topk_weights,
+            router_logits.shape[1],
+            random=simulate_uniform_experts,
             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