Fix MoE reduce-scatterv eligibility check (#32663)
This commit is contained in:
@@ -447,12 +447,18 @@ def should_use_dp_reduce_scatterv():
|
|||||||
Use reduce_scatterv in the standard dispatcher's combine() for DP attention
|
Use reduce_scatterv in the standard dispatcher's combine() for DP attention
|
||||||
with EP, replacing the default all-reduce + dp_scatter path.
|
with EP, replacing the default all-reduce + dp_scatter path.
|
||||||
Only changes the combine (post-kernel) communication; dispatch is unchanged.
|
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 (
|
return (
|
||||||
not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
||||||
and get_moe_a2a_backend().is_none()
|
and get_moe_a2a_backend().is_none()
|
||||||
and is_dp_attention_enabled()
|
and is_dp_attention_enabled()
|
||||||
and get_parallel().attn_dp_size > 1
|
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
|
and get_parallel().moe_ep_size == get_parallel().attn_dp_size
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -894,6 +894,23 @@ class TestForwardFlags(_IsolatedServerArgs):
|
|||||||
self.assertTrue(fwd.flashinfer_trtllm_bypass)
|
self.assertTrue(fwd.flashinfer_trtllm_bypass)
|
||||||
self.assertFalse(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):
|
class TestPublishLifecycle(_IsolatedServerArgs):
|
||||||
"""Publish installs the resolved server_args and seeds the capture tier."""
|
"""Publish installs the resolved server_args and seeds the capture tier."""
|
||||||
|
|||||||
Reference in New Issue
Block a user