Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions python/sglang/srt/layers/moe/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)

Expand Down
17 changes: 17 additions & 0 deletions test/registered/unit/test_runtime_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading