From aa06f6f8a7eedd8d0758b5af3c7c40d06333ef5f Mon Sep 17 00:00:00 2001 From: Sam Li Date: Tue, 28 Jul 2026 07:29:30 -0700 Subject: [PATCH] fix(moe): validate reduce-scatterv group size --- python/sglang/srt/layers/moe/utils.py | 6 ++++++ test/registered/unit/test_runtime_context.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 0f3f48f1e4c8..2287e35949e9 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -445,12 +445,18 @@ def should_use_dp_reduce_scatterv(): Use reduce_scatterv in the standard dispatcher's combine() for DP attention with EP, replacing the default all-reduce + dp_scatter path. Only changes the combine (post-kernel) communication; dispatch is unchanged. + + The reduce_scatterv group is the global TP group, while its variable split + sizes are one entry per attention-DP rank. Therefore this optimization is + valid only when each attention-DP shard has a single rank (attention TP=1). + Configurations with partial attention TP fall back to all-reduce + dp_scatter. """ return ( not should_use_flashinfer_cutlass_moe_fp4_allgather() and get_moe_a2a_backend().is_none() and is_dp_attention_enabled() and get_parallel().attn_dp_size > 1 + and get_parallel().tp_size == get_parallel().attn_dp_size and get_parallel().moe_ep_size == get_parallel().attn_dp_size ) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 35ea11e3f591..0eecb09e59c8 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -894,6 +894,23 @@ def test_mlp_comm_forward_flags(self): self.assertTrue(fwd.flashinfer_trtllm_bypass) self.assertFalse(fwd.flashinfer_trtllm_bypass) + def test_dp_reduce_scatterv_requires_single_rank_attention_dp_shards(self): + from sglang.srt.layers.moe.utils import should_use_dp_reduce_scatterv + + reset_context() + with patch( + "sglang.srt.layers.moe.utils.is_dp_attention_enabled", + return_value=True, + ): + # The optimized path is valid when the collective group and the + # variable-split list have the same number of entries. + with get_parallel().override(tp_size=8, attn_dp_size=8, moe_ep_size=8): + self.assertTrue(should_use_dp_reduce_scatterv()) + + # Otherwise the standard all-reduce plus scatter path must be used. + with get_parallel().override(tp_size=8, attn_dp_size=2, moe_ep_size=2): + self.assertFalse(should_use_dp_reduce_scatterv()) + class TestPublishLifecycle(_IsolatedServerArgs): """Publish installs the resolved server_args and seeds the capture tier."""