From 06ccaef24afeb93293356e0227150b94da5a58b7 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:10:04 -0700 Subject: [PATCH] Fix silently wrong EPLB output with --moe-a2a-backend none (rank-invariant dispatch) (#32962) --- python/sglang/srt/eplb/expert_location.py | 7 +- .../srt/eplb/expert_location_dispatch.py | 42 ++++++- python/sglang/srt/server_args.py | 21 +++- test/registered/ep/test_eplb_no_a2a.py | 110 ++++++++++++++++++ ...e_logical_to_rank_dispatch_physical_map.py | 15 ++- 5 files changed, 183 insertions(+), 12 deletions(-) create mode 100644 test/registered/ep/test_eplb_no_a2a.py diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index 7602fdb1f..0d633f643 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/python/sglang/srt/eplb/expert_location.py @@ -560,8 +560,11 @@ def _compute_logical_to_all_physical_map( physical_expert_id ) - # Replace by the physical expert on local GPU or node if possible - if moe_ep_rank is not None: + # Replace by the physical expert on local GPU or node if possible. Skipped + # without an a2a backend, where all EP ranks must agree on the pick: this + # collapse is per-rank, and the full candidate list is what lets the dispatch + # spread a hot expert over its replicas. See ExpertLocationDispatchInfo. + if moe_ep_rank is not None and server_args.moe_a2a_backend != "none": num_local_gpu_physical_experts = num_physical_experts // ep_size prefer_same_node = _prefer_same_node_experts(server_args) num_gpus_per_node = ( diff --git a/python/sglang/srt/eplb/expert_location_dispatch.py b/python/sglang/srt/eplb/expert_location_dispatch.py index aa144c644..60b01f29e 100644 --- a/python/sglang/srt/eplb/expert_location_dispatch.py +++ b/python/sglang/srt/eplb/expert_location_dispatch.py @@ -23,7 +23,7 @@ from sglang.srt.runtime_context import get_server_args @dataclass class ExpertLocationDispatchInfo: - ep_dispatch_algorithm: Literal["static", "random"] + ep_dispatch_algorithm: Literal["static", "dynamic", "fake", "lp"] # (num_logical_experts,) partial_logical_to_rank_dispatch_physical_map: Optional[torch.Tensor] # (num_logical_experts, X) @@ -31,10 +31,17 @@ class ExpertLocationDispatchInfo: # (num_logical_experts,) partial_logical_to_all_physical_map_num_valid: torch.Tensor num_physical_experts: int + # Whether every rank must pick the same physical expert for a token. True + # without an a2a backend: all EP ranks then run the MoE over the same tokens + # and sum their partial outputs, so a rank-dependent pick counts a replicated + # logical expert several times. With one, each rank dispatches only its own + # tokens and is free to disagree. + rank_invariant: bool = False @classmethod def init_new(cls, layer_id: int): - ep_dispatch_algorithm = get_server_args().ep_dispatch_algorithm + server_args = get_server_args() + ep_dispatch_algorithm = server_args.ep_dispatch_algorithm expert_location_metadata = get_global_expert_location_metadata() assert expert_location_metadata is not None @@ -43,6 +50,7 @@ class ExpertLocationDispatchInfo: return cls( ep_dispatch_algorithm=ep_dispatch_algorithm, + rank_invariant=server_args.moe_a2a_backend == "none", partial_logical_to_rank_dispatch_physical_map=( expert_location_metadata.logical_to_rank_dispatch_physical_map[ layer_id, : @@ -107,15 +115,37 @@ def _topk_ids_logical_to_physical_static( def _topk_ids_logical_to_physical_dynamic( topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo] ) -> torch.Tensor: + """Spread each (token, logical expert) over that logical expert's replicas. + + Under ``rank_invariant`` the replica is picked by token row rather than at + random: ``torch.randint`` reads the default CUDA generator, so agreement + across ranks would rest on their philox offsets staying aligned, which + nothing asserts, and it would make greedy requests non-reproducible. Both + pick evenly, which is the load split EPLB's placement solver assumes. + + Row indexing leaves replicas unused for a single-row batch, where there is no + imbalance to fix anyway. + """ topk_ids_original_shape = topk_ids.shape original_dtype = topk_ids.dtype device = topk_ids.device topk_ids = topk_ids.flatten() - chosen_dispatch_index = ( - torch.randint(0, 65536, topk_ids.shape, dtype=torch.int32, device=device) - % info.partial_logical_to_all_physical_map_num_valid[topk_ids] - ) + num_valid = info.partial_logical_to_all_physical_map_num_valid[topk_ids] + if info.rank_invariant: + slots_per_token = ( + topk_ids_original_shape[-1] if len(topk_ids_original_shape) > 1 else 1 + ) + row_index = ( + torch.arange(topk_ids.shape[0], dtype=num_valid.dtype, device=device) + // slots_per_token + ) + chosen_dispatch_index = row_index % num_valid + else: + chosen_dispatch_index = ( + torch.randint(0, 65536, topk_ids.shape, dtype=torch.int32, device=device) + % num_valid + ) topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index] if topk_ids.dtype != original_dtype: topk_ids = topk_ids.to(original_dtype) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2740aa872..1f874fa5d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6710,10 +6710,29 @@ class ServerArgs: "EPLB is enabled. The expert_distribution_recorder_mode is automatically set." ) + # Without an a2a backend all EP ranks run the MoE over the same tokens and + # sum their partial outputs, so the pick has to agree across ranks. + needs_rank_invariant_dispatch = self._resolved().moe_a2a_backend == "none" + if (self.enable_eplb or (self.init_expert_location != "trivial")) and ( self.ep_dispatch_algorithm is None ): - self.ep_dispatch_algorithm = "static" + self.ep_dispatch_algorithm = ( + "dynamic" if needs_rank_invariant_dispatch else "static" + ) + + # `dynamic` / `fake` switch to the row-index pick; `static` reads a + # per-rank table and `lp` samples inside its kernel. + if needs_rank_invariant_dispatch and self.ep_dispatch_algorithm in ( + "static", + "lp", + ): + raise ValueError( + f"--ep-dispatch-algorithm {self.ep_dispatch_algorithm} picks a " + "different physical replica per rank, which only holds up when an " + "a2a backend routes each token to a single rank. Use " + "--ep-dispatch-algorithm dynamic with --moe-a2a-backend none." + ) if self.enable_eplb and self.ep_join_mode != "scale": assert self._resolved().ep_size > 1 diff --git a/test/registered/ep/test_eplb_no_a2a.py b/test/registered/ep/test_eplb_no_a2a.py new file mode 100644 index 000000000..1a9e5d720 --- /dev/null +++ b/test/registered/ep/test_eplb_no_a2a.py @@ -0,0 +1,110 @@ +"""EPLB with redundant experts on the no-a2a MoE path (--moe-a2a-backend none). + +There, all EP ranks run the MoE over the same tokens and sum their partial +outputs, so the logical->physical pick has to be identical on every rank -- a +rank-dependent one counts a replicated logical expert several times and silently +degrades output instead of failing. + +`test/manual/ep/test_eplb.py` covers EPLB with an a2a backend. +""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, + try_cached_model, +) + +register_cuda_ci(est_time=420, suite="nightly-eval-text-2-gpu", nightly=True) + +# 72 routed experts + 48 replicas = 120 physical, 60 per rank, so two thirds of +# the routed (token, expert) pairs get double-counted when ranks disagree. At 24 +# replicas the score only fell to 0.575 against the 0.60 threshold. +NUM_REDUNDANT_EXPERTS = 48 + + +class TestEPLBNoA2A(CustomTestCase): + """Initial placement, no rebalance during the eval. + + Guards the candidate-map half of the fix: only the initial placement goes + through `_compute_logical_to_all_physical_map`, where the rank-local collapse + lives, so a regression there is invisible once rebalancing starts. + """ + + extra_args = [] + # Never reached by a 200-question eval. Also sizes the expert-distribution + # recorder buffer, so it cannot be made arbitrarily large. + rebalance_num_iterations = "20000" + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(DEFAULT_MODEL_NAME_FOR_TEST_MLA) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "2", + "--ep-size", + "2", + "--enable-eplb", + "--ep-num-redundant-experts", + str(NUM_REDUNDANT_EXPERTS), + "--eplb-rebalance-num-iterations", + cls.rebalance_num_iterations, + "--mem-fraction-static", + "0.5", + *cls.extra_args, + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(metrics) + + # Measured over six runs: 0.625-0.66 correct, 0.415-0.435 with the + # rank-dependent pick restored. 0.60 (what the other EP tests use for + # this model) sits under a sigma of the low end, so leave room. + self.assertGreater(metrics["score"], 0.55) + + +class TestEPLBNoA2ADPAttention(TestEPLBNoA2A): + """DP attention -- the MoE runs over the DP-gathered global token buffer -- + plus ~70 real rebalances, which exercise the post-rebalance placements and + the expert-weight migration.""" + + extra_args = [ + "--enable-dp-attention", + "--dp", + "2", + ] + rebalance_num_iterations = "50" + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/eplb/test_compute_logical_to_rank_dispatch_physical_map.py b/test/registered/unit/eplb/test_compute_logical_to_rank_dispatch_physical_map.py index ea73a7e0d..fa2c729fb 100644 --- a/test/registered/unit/eplb/test_compute_logical_to_rank_dispatch_physical_map.py +++ b/test/registered/unit/eplb/test_compute_logical_to_rank_dispatch_physical_map.py @@ -17,9 +17,18 @@ from sglang.srt.eplb.expert_location import ( from sglang.test.test_utils import CustomTestCase -def _make_server_args(ep_size: int, nnodes: int): - """Minimal server_args stub for expert placement tests.""" - return types.SimpleNamespace(ep_size=ep_size, nnodes=nnodes, ep_join_mode=None) +def _make_server_args(ep_size: int, nnodes: int, moe_a2a_backend: str = "deepep"): + """Minimal server_args stub for expert placement tests. + + `moe_a2a_backend` defaults to an a2a backend because these tests cover the + rank-local collapse, which is skipped when there is no a2a backend. + """ + return types.SimpleNamespace( + ep_size=ep_size, + nnodes=nnodes, + ep_join_mode=None, + moe_a2a_backend=moe_a2a_backend, + ) def _make_logical_to_all_physical_map(