Support moe_dp_size = 1 for various attention_cp_size (#22003)

Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>

python/sglang/srt/distributed/parallel_state.py

@@ -1870,6 +1870,7 @@ def initialize_model_parallel(
                 )
                 ranks = list(range(st, en))
                 group_ranks.append(ranks)
+
         _ATTN_TP = init_model_parallel_group(
             group_ranks,
             get_world_group().local_rank,
@@ -1888,8 +1889,12 @@ def initialize_model_parallel(
 
     global _MOE_DP
     assert _MOE_DP is None, "moe data parallel group is already initialized"
-    # gpus_per_pp_stage = tensor_model_parallel_size * attention_context_model_parallel_size
-    if moe_dp_size == tensor_model_parallel_size:
+    if attn_cp_size > moe_dp_size:
+        # When moe_dp_size < attn_cp_size, CP ranks must share tokens before MoE.
+        # The MOE_DP group includes these CP partners, so the existing DP
+        # allgather/scatter handles the token sharing.
+        _MOE_DP = _ATTN_CP
+    elif moe_dp_size == tensor_model_parallel_size:
         _MOE_DP = _TP
     else:
         group_ranks = []
@@ -2204,6 +2209,12 @@ def destroy_model_parallel():
     _MOE_TP = None
 
     global _ATTN_CP
+    global _MOE_DP
+    # Destroy _MOE_DP before _ATTN_CP since it may alias _ATTN_CP.
+    # Only destroy if not aliasing another group.
+    if _MOE_DP and _MOE_DP is not _ATTN_CP and _MOE_DP is not _TP:
+        _MOE_DP.destroy()
+    _MOE_DP = None
     if _ATTN_CP:
         _ATTN_CP.destroy()
     _ATTN_CP = None
@@ -2213,11 +2224,6 @@ def destroy_model_parallel():
         _ATTN_TP.destroy()
     _ATTN_TP = None
 
-    global _MOE_DP
-    if _MOE_DP:
-        _MOE_DP.destroy()
-    _MOE_DP = None
-
     global _PDMUX_PREFILL_TP_GROUP
     if _PDMUX_PREFILL_TP_GROUP:  # type: ignore[union-attr]
         _PDMUX_PREFILL_TP_GROUP.destroy()

python/sglang/srt/layers/communicator.py

@@ -50,8 +50,12 @@ from sglang.srt.layers.dp_attention import (
     get_dp_global_num_tokens,
     get_global_dp_buffer,
     get_local_dp_buffer,
+    get_moe_cp_rank,
+    get_moe_cp_size,
     is_allocation_symmetric,
     is_dp_attention_enabled,
+    is_enable_moe_cp_allgather,
+    moe_cp_all_gather_into_tensor,
 )
 from sglang.srt.layers.flashinfer_comm_fusion import is_flashinfer_allreduce_unavailable
 from sglang.srt.layers.moe import (
@@ -185,11 +189,13 @@ class ScatterMode(Enum):
     SCATTERED: [a, b, c, d]
     TP_ATTN_FULL: [ab, ab, cd, cd], i.e. all ranks inside a TP attn group have full data of the group
     FULL: [abcd, abcd, abcd, abcd]
+    MOE_FULL: full within the MoE group (cp_per_moe CP chunks), used when moe_dp_size < attn_cp_size
     """
 
     SCATTERED = auto()
     TP_ATTN_FULL = auto()
     FULL = auto()
+    MOE_FULL = auto()
 
     @staticmethod
     def model_input_output():
@@ -362,15 +368,16 @@ class LayerScatterModes:
     @classmethod
     def _compute_mlp_mode(cls, context: _LayerModeComputationContext):
         if context.is_layer_sparse:
-            return (
-                ScatterMode.SCATTERED
+            if (
                 # Token dispatch/combine will be handled outside of LayerCommunicator for these modes.
-                if (
-                    not get_moe_a2a_backend().is_none()
-                    or should_use_flashinfer_cutlass_moe_fp4_allgather()
-                )
-                else ScatterMode.FULL
-            )
+                not get_moe_a2a_backend().is_none()
+                or should_use_flashinfer_cutlass_moe_fp4_allgather()
+            ):
+                return ScatterMode.SCATTERED
+            # NSA CP doesn't support MOE_FULL yet; fall back to FULL
+            if is_enable_moe_cp_allgather() and not is_nsa_enable_prefill_cp():
+                return ScatterMode.MOE_FULL
+            return ScatterMode.FULL
         else:
             return (
                 ScatterMode.SCATTERED
@@ -392,7 +399,7 @@ class LayerScatterModes:
         mlp_mode = cls._compute_mlp_mode(context)
         if mlp_mode == ScatterMode.SCATTERED:
             return ScatterMode.SCATTERED
-        if mlp_mode == ScatterMode.FULL:
+        if mlp_mode in (ScatterMode.FULL, ScatterMode.MOE_FULL):
             return ScatterMode.TP_ATTN_FULL
         raise NotImplementedError
 
@@ -405,7 +412,7 @@ class LayerScatterModes:
             if cls._should_gather_for_tbo(context):
                 return ScatterMode.TP_ATTN_FULL
             return ScatterMode.SCATTERED
-        if mlp_mode == ScatterMode.FULL:
+        if mlp_mode in (ScatterMode.FULL, ScatterMode.MOE_FULL):
             return ScatterMode.TP_ATTN_FULL
         raise NotImplementedError
 
@@ -705,6 +712,13 @@ class LayerCommunicator:
     def should_fuse_mlp_allreduce_with_next_layer(
         self, forward_batch: ForwardBatch
     ) -> bool:
+        # When MOE_FULL is active (moe_cp allgather), fusion must be disabled because
+        # 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():
+            return False
+
         if (
             is_dp_attention_enabled()
             and self._speculative_algo is not None
@@ -760,6 +774,7 @@ class CommunicateContext:
         attn_cp_rank = get_attention_cp_rank()
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
+        moe_cp_size = get_moe_cp_size()
         process_group_sizes = {
             ScatterMode.SCATTERED: 1,
             ScatterMode.TP_ATTN_FULL: attn_tp_size,
@@ -767,6 +782,7 @@ class CommunicateContext:
             # With context parallel enabled, we should exclude
             # the attn_cp_size from the total tp_size
             ScatterMode.FULL: tp_size // attn_cp_size,
+            ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size),
         }
         return cls(
             process_group_sizes=process_group_sizes,
@@ -883,6 +899,19 @@ class CommunicateWithAllReduceAndLayerNormFn:
                 residual_input_mode=residual_input_mode,
             )
 
+        if (
+            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
+            and (
+                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
+            )
+            and (hidden_states_output_mode == ScatterMode.MOE_FULL)
+            and (residual_output_mode == ScatterMode.TP_ATTN_FULL)
+        ):
+            return partial(
+                CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual_moe,
+                residual_input_mode=residual_input_mode,
+            )
+
         if (
             (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
             and (
@@ -1016,6 +1045,77 @@ class CommunicateWithAllReduceAndLayerNormFn:
         hidden_states = layernorm(residual)
         return hidden_states, residual
 
+    @staticmethod
+    def _gather_hidden_states_and_residual_moe(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch,
+        layernorm: torch.nn.Module,
+        context: CommunicateContext,
+        *,
+        residual_input_mode,
+    ):
+        """Allgather tokens for MoE when moe_dp_size < attn_cp_size.
+
+        Steps:
+          1. Standard attn-TP all-reduce + optional DP allgather + layernorm (same as
+             _gather_hidden_states_and_residual for the dp>1 case, or simple all-reduce
+             + layernorm for dp==1).
+          2. moe_cp allgather: gather tokens from cp_per_moe CP ranks so each rank holds
+             all tokens for its MoE group.
+
+        Residual is left at TP_ATTN_FULL throughout.
+        """
+        # Early return on empty tensor is safe for MOE_CP because:
+        # - During CP extend: zigzag split guarantees all CP ranks have non-zero tokens,
+        #   so no rank hits this path while others proceed to the allgather.
+        # - During decode: moe_cp allgather is skipped (guarded by is_context_parallel_extend).
+        # - CUDA graph warmup: not applicable when --disable-piecewise-cuda-graph is used.
+        if hidden_states.shape[0] == 0:
+            return hidden_states, residual
+
+        # Step 1: Standard all-reduce/DP-allgather + layernorm (reuse existing logic).
+        hidden_states, residual = (
+            CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual(
+                hidden_states=hidden_states,
+                residual=residual,
+                forward_batch=forward_batch,
+                layernorm=layernorm,
+                context=context,
+                residual_input_mode=residual_input_mode,
+            )
+        )
+
+        # Step 2: moe_cp allgather — gather across cp_per_moe CP ranks.
+        # Only active during prefill (context-parallel extend); decode keeps existing path.
+        moe_cp_size = get_moe_cp_size()
+        if (
+            moe_cp_size > 1
+            and hidden_states.shape[0] > 0
+            and forward_batch.forward_mode.is_context_parallel_extend()
+            and forward_batch.attn_cp_metadata is not None
+        ):
+            # Zigzag split can produce unequal token counts across CP ranks
+            # (when seq_len % (cp_size * 2) != 0). NCCL allgather requires
+            # equal input sizes, so pad to the max per-rank token count.
+            per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
+            max_tokens = max(per_rank_tokens)
+            pad_size = max_tokens - hidden_states.shape[0]
+            if pad_size > 0:
+                hidden_states = torch.nn.functional.pad(
+                    hidden_states, [0, 0, 0, pad_size]
+                )
+
+            output = torch.empty(
+                (max_tokens * moe_cp_size, hidden_states.shape[1]),
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+            moe_cp_all_gather_into_tensor(output, hidden_states)
+            hidden_states = output
+
+        return hidden_states, residual
+
 
 class CommunicateSummableTensorPairFn:
     """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed."""
@@ -1069,6 +1169,13 @@ class CommunicateSummableTensorPairFn:
         ):
             return CommunicateSummableTensorPairFn._scatter
 
+        if (
+            (hidden_states_input_mode == ScatterMode.MOE_FULL)
+            and (residual_input_mode == ScatterMode.TP_ATTN_FULL)
+            and (output_mode == ScatterMode.TP_ATTN_FULL)
+        ):
+            return CommunicateSummableTensorPairFn._scatter_hidden_states_moe
+
         raise NotImplementedError(
             f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}"
         )
@@ -1138,3 +1245,50 @@ class CommunicateSummableTensorPairFn:
         tensor_list = list(hidden_states.tensor_split(context.attn_tp_size))
         hidden_states = tensor_list[context.attn_tp_rank]
         return hidden_states, residual
+
+    @staticmethod
+    def _scatter_hidden_states_moe(
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        forward_batch: ForwardBatch,
+        context: CommunicateContext,
+        **kwargs,
+    ):
+        """Scatter MoE output back to TP_ATTN_FULL after MOE_FULL computation.
+
+        After moe_tensor_model_parallel_all_reduce (which runs unconditionally since
+        use_reduce_scatter=False for this path), all ranks in the moe_cp group hold the
+        full MoE result for all cp_per_moe token chunks. We simply slice out this rank's
+        CP-local portion.
+
+        If DP>1, further scatter back to the local DP slice.
+        """
+        # Only scatter back during prefill; decode was never allgathered so no-op.
+        # Safe w.r.t. empty tensors: same reasoning as _gather_hidden_states_and_residual_moe
+        # — CP extend always has non-zero tokens per rank, and decode skips this path.
+        moe_cp_size = get_moe_cp_size()
+        if (
+            moe_cp_size > 1
+            and forward_batch.forward_mode.is_context_parallel_extend()
+            and forward_batch.attn_cp_metadata is not None
+        ):
+            moe_cp_rank = get_moe_cp_rank()
+            # The allgather was padded to max_tokens_per_rank (equal chunks).
+            # Extract this rank's actual (non-padded) tokens from its chunk.
+            per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
+            max_tokens_per_rank = max(per_rank_tokens)
+            actual_local_tokens = per_rank_tokens[moe_cp_rank]
+            hidden_states = hidden_states.narrow(
+                0, moe_cp_rank * max_tokens_per_rank, actual_local_tokens
+            ).contiguous()
+
+        # DP scatter (if DP attention is enabled)
+        if context.attn_dp_size > 1:
+            hidden_states_output, global_hidden_states = (
+                get_local_dp_buffer(),
+                hidden_states,
+            )
+            dp_scatter(hidden_states_output, global_hidden_states, forward_batch)
+            hidden_states = hidden_states_output
+
+        return hidden_states, residual

python/sglang/srt/layers/dp_attention.py

@@ -18,6 +18,9 @@ from sglang.srt.distributed import (
     get_attn_tensor_model_parallel_rank,
     get_attn_tensor_model_parallel_world_size,
     get_attn_tp_group,
+)
+from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
+from sglang.srt.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
     get_tp_group,
@@ -580,5 +583,30 @@ def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
     return get_attention_cp_group().all_gather_into_tensor(output, input)
 
 
+def get_moe_cp_group() -> GroupCoordinator:
+    """Returns the MOE_DP group, which includes CP partners when attn_cp_size > moe_dp_size."""
+    return _get_moe_dp_group()
+
+
+def get_moe_cp_rank() -> int:
+    return _get_moe_dp_group().rank_in_group
+
+
+def get_moe_cp_size() -> int:
+    return _get_moe_dp_group().world_size
+
+
+def is_enable_moe_cp_allgather() -> bool:
+    """True when moe_dp_size < attn_cp_size, requiring allgather across CP ranks before MoE."""
+    from sglang.srt.server_args import get_global_server_args
+
+    sa = get_global_server_args()
+    return sa.attn_cp_size > sa.moe_dp_size
+
+
+def moe_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
+    return _get_moe_dp_group().all_gather_into_tensor(output, input)
+
+
 def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor):
     return get_attention_tp_group().all_gather(input, output_tensor_list=output_list)

python/sglang/srt/layers/utils/cp_utils.py

@@ -51,16 +51,14 @@ def is_prefill_cp_in_seq_split():
 
 
 def can_cp_split(seq_len: int, cp_size: int, forward_batch):
-    # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
-    # Note: (self.cp_size * 2) To achieve load balancing for seq computation,
-    # the seq data needs to be divided and recombined at twice the size of cp_size.
+    # CP metadata (zigzag split) only supports batch=1 for now.
     cur_cp_seq_len = seq_len // (cp_size * 2)
-    # print("DEBUG: can_cp_split", cur_cp_seq_len, cp_size, forward_batch.forward_mode.is_context_parallel_extend(), is_prefill_context_parallel_enabled(), flush=True)
     if (
         cur_cp_seq_len != 0
         and cp_size > 1
         and forward_batch.forward_mode.is_context_parallel_extend()
         and is_prefill_context_parallel_enabled()
+        and forward_batch.seq_lens_cpu.shape[0] == 1
     ):
         return True
     else:

python/sglang/srt/models/qwen2_moe.py

@@ -33,6 +33,9 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.distributed.parallel_state import (
+    get_attn_context_model_parallel_world_size,
+)
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
@@ -709,6 +712,7 @@ class Qwen2MoeModel(nn.Module):
         self.pp_group = get_pp_group()
 
         self.moe_dp_size = get_moe_data_parallel_world_size()
+        self.attn_cp_size = get_attn_context_model_parallel_world_size()
 
         if self.pp_group.is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -830,7 +834,7 @@ class Qwen2MoeModel(nn.Module):
         ):
             hidden_states = cp_all_gather_rerange_output(
                 hidden_states,
-                self.moe_dp_size,
+                self.attn_cp_size,
                 forward_batch,
                 torch.cuda.current_stream(),
             )

python/sglang/srt/models/qwen3_moe.py

@@ -968,9 +968,10 @@ class Qwen3MoeForCausalLM(nn.Module):
         self.attn_cp_rank = get_attn_context_model_parallel_rank()
         self.moe_dp_size = get_moe_data_parallel_world_size()
 
-        assert (
-            self.attn_cp_size == self.moe_dp_size
-        ), "Attention context parallel size must be equal to MoE context parallel size"
+        assert self.attn_cp_size % self.moe_dp_size == 0, (
+            f"attn_cp_size ({self.attn_cp_size}) must be divisible by "
+            f"moe_dp_size ({self.moe_dp_size})"
+        )
 
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens

python/sglang/srt/server_args.py

@@ -2780,6 +2780,11 @@ class ServerArgs:
                 not self.enable_aiter_allreduce_fusion
             ), "Aiter allreduce fusion is not supported with context parallelism"
 
+        if self.attn_cp_size != self.moe_dp_size:
+            assert (
+                self.moe_dp_size == 1
+            ), "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1"
+
     def _handle_data_parallelism(self):
         if self.dp_size == 1:
             self.enable_dp_attention = False

test/registered/4-gpu-models/test_qwen3_30b.py

@@ -73,5 +73,60 @@ class TestQwen330B(CustomTestCase):
         self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY)
 
 
+class TestQwen330BCP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = QWEN3_30B_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                "4",
+                "--moe-dp-size",
+                "1",
+                "--ep-size",
+                "4",
+                "--attn-cp-size",
+                "2",
+                "--enable-prefill-context-parallel",
+                "--cuda-graph-max-bs",
+                "32",
+                "--max-running-requests",
+                "32",
+                "--trust-remote-code",
+                "--disable-piecewise-cuda-graph",
+                "--model-loader-extra-config",
+                '{"enable_multithread_load": true, "num_threads": 64}',
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            model=self.model,
+            eval_name="gsm8k",
+            num_shots=5,
+            num_examples=200,
+            max_tokens=16000,
+            num_threads=128,
+            repeat=1,
+            temperature=0.6,
+            top_p=0.95,
+            top_k=20,
+            base_url=self.base_url,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY)
+
+
 if __name__ == "__main__":
     unittest.main()