From 65f79571421651431b15cf921e7b691daee6f596 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:14:31 -0700 Subject: [PATCH] fix: gather CP-sharded tokens before TP-sharded dense MLP under prefill CP (#38078) --- python/sglang/srt/layers/communicator.py | 30 ++++++++++++---- .../test_layer_communicator_fusion_gate.py | 20 ++++++++--- .../test_layer_scatter_modes_cp_dense_mlp.py | 36 +++++++++++++++++++ 3 files changed, 76 insertions(+), 10 deletions(-) create mode 100644 test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index c395f96bf..459394b4f 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -425,11 +425,17 @@ class LayerScatterModes: return ScatterMode.MOE_FULL return ScatterMode.FULL else: - return ( - ScatterMode.SCATTERED - if enable_moe_dense_fully_dp() - else ScatterMode.FULL - ) + if enable_moe_dense_fully_dp(): + return ScatterMode.SCATTERED + # A TP-sharded dense MLP reduces over the whole TP group, which spans + # every CP rank; a CP-sharded prefill must gather tokens across CP + # first or the all-reduce sums different tokens' partial outputs. + # MLA/DSA CP models do this in DSACPLayerCommunicator instead. + if _generic_prefill_cp_shards_tokens() and not ( + is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled() + ): + return ScatterMode.MOE_FULL + return ScatterMode.FULL @classmethod def _should_gather_for_tbo(cls, context: _LayerModeComputationContext): @@ -467,6 +473,15 @@ def enable_moe_dense_fully_dp(): return get_parallel().moe_dense_tp_size == 1 +def _generic_prefill_cp_shards_tokens() -> bool: + """Whether the strategy prefill CP path shards prefill tokens across CP ranks.""" + # Local import: module-level CP helper imports here are circular (#27014). + from sglang.srt.layers.cp.utils import enable_cp_v2 + + parallel = get_parallel() + return parallel.attn_cp_size > 1 and parallel.enable_prefill_cp and enable_cp_v2() + + def enable_dwdp(): return get_parallel().dwdp_size > 1 @@ -890,7 +905,10 @@ class LayerCommunicator: # the fusion path skips postprocess_layer which contains the moe_cp scatter. # Without scatter, hidden_states remain at MOE_FULL size while residual is at # TP_ATTN_FULL size, causing a shape mismatch. - if is_enable_moe_cp_allgather(): + if ( + is_enable_moe_cp_allgather() + or self.layer_scatter_modes.mlp_mode == ScatterMode.MOE_FULL + ): return False # Fusing makes the next layer's residual+LN absorb the post-experts diff --git a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py index 16f49cf15..6241aa9e5 100644 --- a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py +++ b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py @@ -11,10 +11,10 @@ from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def _fake_communicator(): +def _fake_communicator(mlp_mode=ScatterMode.TP_ATTN_FULL): return types.SimpleNamespace( _speculative_algo=None, - layer_scatter_modes=types.SimpleNamespace(mlp_mode=ScatterMode.TP_ATTN_FULL), + layer_scatter_modes=types.SimpleNamespace(mlp_mode=mlp_mode), is_last_layer=False, _context=types.SimpleNamespace(tp_size=4), ) @@ -31,7 +31,9 @@ class TestFuseMlpAllReduceGate(CustomTestCase): Qwen3-30B-A3B with --tp-size 4 --ep-size 2. """ - def _should_fuse(self, *, moe_ep_size, moe_tp_size): + def _should_fuse( + self, *, moe_ep_size, moe_tp_size, mlp_mode=ScatterMode.TP_ATTN_FULL + ): forward_batch = types.SimpleNamespace( input_ids=types.SimpleNamespace(shape=(8,)) ) @@ -48,7 +50,7 @@ class TestFuseMlpAllReduceGate(CustomTestCase): ), ): return LayerCommunicator.should_fuse_mlp_allreduce_with_next_layer( - _fake_communicator(), forward_batch + _fake_communicator(mlp_mode), forward_batch ) def test_hybrid_ep_tp_does_not_fuse(self): @@ -60,6 +62,16 @@ class TestFuseMlpAllReduceGate(CustomTestCase): def test_pure_ep_still_fuses(self): self.assertTrue(self._should_fuse(moe_ep_size=4, moe_tp_size=1)) + def test_moe_full_layer_does_not_fuse(self): + # Fusion skips postprocess_layer, which holds the CP scatter; a dense + # MOE_FULL layer (moe_dp_size == attn_cp_size) is not caught by the + # is_enable_moe_cp_allgather gate. + self.assertFalse( + self._should_fuse( + moe_ep_size=1, moe_tp_size=4, mlp_mode=ScatterMode.MOE_FULL + ) + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py b/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py new file mode 100644 index 000000000..4d1867788 --- /dev/null +++ b/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py @@ -0,0 +1,36 @@ +import unittest +from unittest.mock import patch + +from sglang.srt.layers import communicator as comm +from sglang.srt.layers.communicator import LayerScatterModes, ScatterMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestDenseMlpScatterModeUnderPrefillCP(CustomTestCase): + """A TP-sharded dense MLP under prefill CP must gather tokens across CP ranks + before its all-reduce, or CP pairs sum partial outputs of different tokens + (issue #38019: Qwen3-32B emitted garbage that never reached EOS).""" + + def test_dense_mlp_gathers_across_cp(self): + with ( + patch.object(comm, "_generic_prefill_cp_shards_tokens", return_value=True), + patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False), + patch.object(comm, "is_mla_prefill_cp_enabled", return_value=False), + patch.object(comm, "enable_moe_dense_fully_dp", return_value=False), + ): + modes = LayerScatterModes.init_new( + layer_id=1, + num_layers=4, + is_layer_sparse=False, + is_previous_layer_sparse=False, + is_next_layer_sparse=False, + ) + self.assertEqual(modes.mlp_mode, ScatterMode.MOE_FULL) + self.assertEqual(modes.layer_output_mode, ScatterMode.TP_ATTN_FULL) + + +if __name__ == "__main__": + unittest.main()