diff --git a/python/sglang/srt/layers/moe/__init__.py b/python/sglang/srt/layers/moe/__init__.py index 7086c8bdc..4dc3992f2 100644 --- a/python/sglang/srt/layers/moe/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 6e047ac88..b429271e7 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 6638fe83b..3f808ee59 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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 diff --git a/test/manual/layers/moe/test_simulate_balanced_routing.py b/test/manual/layers/moe/test_simulate_balanced_routing.py new file mode 100644 index 000000000..314e44525 --- /dev/null +++ b/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()