minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: jonnykong <jonnykong@fb.com>
This commit is contained in:
Jonny Kong
2026-08-28 14:22:45 -07:00
committed by GitHub
co-authored by jonnykong
parent f65b2b2b15
commit 4d78d59e51
4 changed files with 309 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/__init__.py
+2
View File
@@ -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
+133 -42
View File
@@ -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
+9
View File
@@ -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
+165
View File
@@ -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()