fix: gather CP-sharded tokens before TP-sharded dense MLP under prefill CP (#38078)

python/sglang/srt/layers/communicator.py

@@ -425,11 +425,17 @@ class LayerScatterModes:
                 return ScatterMode.MOE_FULL
             return ScatterMode.FULL
         else:
-            return (
+            if enable_moe_dense_fully_dp():
-                ScatterMode.SCATTERED
+                return ScatterMode.SCATTERED
-                if enable_moe_dense_fully_dp()
+            # A TP-sharded dense MLP reduces over the whole TP group, which spans
-                else ScatterMode.FULL
+            # 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

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()

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()