From 3dc1491c95b48fb2950a469401f24a68764c9fd6 Mon Sep 17 00:00:00 2001 From: Shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 21 Apr 2026 02:58:19 +0800 Subject: [PATCH] Support moe_dp_size = 1 for various attention_cp_size (#22003) Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co> --- .../sglang/srt/distributed/parallel_state.py | 20 +- python/sglang/srt/layers/communicator.py | 174 +++++++++++++++++- python/sglang/srt/layers/dp_attention.py | 28 +++ python/sglang/srt/layers/utils/cp_utils.py | 6 +- python/sglang/srt/models/qwen2_moe.py | 6 +- python/sglang/srt/models/qwen3_moe.py | 7 +- python/sglang/srt/server_args.py | 5 + .../registered/4-gpu-models/test_qwen3_30b.py | 55 ++++++ 8 files changed, 276 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index dad005fc2..b4f4aed06 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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() diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 0b0953102..d0096f4d0 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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 diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index f5fcd1757..09d307aeb 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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) diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 40aeb750b..b6981931f 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/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: diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 4cafaab6e..f01632ac7 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/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(), ) diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index cfd52b654..b94ba030c 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1b62823e3..550bc1dd5 100644 --- a/python/sglang/srt/server_args.py +++ b/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 diff --git a/test/registered/4-gpu-models/test_qwen3_30b.py b/test/registered/4-gpu-models/test_qwen3_30b.py index 079dbc985..487445a24 100644 --- a/test/registered/4-gpu-models/test_qwen3_30b.py +++ b/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()