[DSV4] Support megamoe for CP (#29569)

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -11,6 +11,98 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def validate_deepseek_v4_mega_moe_token_budget(
+    server_args: ServerArgs,
+) -> None:
+    """Ensure the DSV4 prefill budget fits MegaMoE's per-rank buffer."""
+    mega_moe_enabled = (
+        server_args.moe_a2a_backend == "megamoe"
+        or envs.SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE.get()
+    )
+    if not mega_moe_enabled or server_args.disaggregation_mode == "decode":
+        # decode node will skip the check because decode bs is not relevant with --chunk-prefill-size
+        return
+
+    if server_args.pp_size > 1 and server_args.enable_dynamic_chunking:
+        return
+
+    if (
+        server_args.chunked_prefill_size is None
+        or server_args.chunked_prefill_size <= 0
+    ):
+        raise ValueError(
+            "DeepSeekV4 with MegaMoE requires chunked prefill to be enabled. "
+            "Set --chunked-prefill-size to a positive value; "
+            "--chunked-prefill-size=-1 is unsafe because MegaMoE's per-rank "
+            "token requirement would not have a strict prefill-forward bound."
+        )
+
+    if server_args.enable_prefill_cp:
+        token_partition_size = server_args.attn_cp_size
+        token_partition_name = "attn_cp_size"
+        token_alignment = 1
+        local_chunked_prefill_size = (
+            server_args.chunked_prefill_size + token_partition_size - 1
+        ) // token_partition_size
+    elif server_args.enable_dp_attention:
+        token_partition_size = server_args.dp_size
+        token_partition_name = "dp_size"
+        token_alignment = max(
+            server_args.tp_size // server_args.dp_size // server_args.attn_cp_size,
+            1,
+        )
+        local_chunked_prefill_size = (
+            server_args.chunked_prefill_size // token_partition_size
+        )
+    else:
+        # Pure TP and PP with static chunking are handled here.
+        token_partition_size = 1
+        token_partition_name = "none"
+        # global_num_tokens will ceil_align to attn_tp_size so the validation needs to do alignment as well
+        token_alignment = max(
+            server_args.tp_size // token_partition_size // server_args.attn_cp_size,
+            1,
+        )
+        local_chunked_prefill_size = server_args.chunked_prefill_size
+
+    if local_chunked_prefill_size <= 0:
+        raise ValueError(
+            "DeepSeekV4 with MegaMoE requires a positive effective per-rank "
+            "chunked prefill size. "
+            f"Current values: chunked_prefill_size="
+            f"{server_args.chunked_prefill_size}, "
+            f"token_partition={token_partition_name}, "
+            f"token_partition_size={token_partition_size}."
+        )
+
+    required_tokens_per_rank = (
+        (local_chunked_prefill_size + token_alignment - 1)
+        // token_alignment
+        * token_alignment
+    )
+    max_tokens_per_rank = (
+        envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.get()
+    )
+    if max_tokens_per_rank < required_tokens_per_rank:
+        raise ValueError(
+            "DeepSeekV4 with MegaMoE requires "
+            "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to "
+            "cover each rank's effective prefill token budget. "
+            f"Current values: chunked_prefill_size="
+            f"{server_args.chunked_prefill_size}, "
+            f"token_partition={token_partition_name}, "
+            f"token_partition_size={token_partition_size}, "
+            f"token_alignment={token_alignment}, "
+            f"required_per_rank={required_tokens_per_rank}, "
+            "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK="
+            f"{max_tokens_per_rank}. Set "
+            "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK to at "
+            f"least {required_tokens_per_rank}, or lower "
+            "--chunked-prefill-size until the effective per-rank budget fits. "
+            "Otherwise MegaMoE falls back to the fused MoE path at runtime."
+        )
+
+
 def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None:
     """Residual imperative arm of the DeepSeek V4 defaults.
 
@@ -85,6 +177,12 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
     assert (
         server_args.tp_size <= 8
     ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues."
+    if server_args.moe_a2a_backend not in ("none", "deepep", "megamoe"):
+        raise ValueError(
+            "DeepSeekV4 CP supports moe_a2a_backend in "
+            "('none', 'deepep', 'megamoe'), "
+            f"got {server_args.moe_a2a_backend!r}."
+        )
     logger.warning(
         "Disabling SGLANG_OPT_FLASHMLA_SPARSE_PREFILL because DeepSeekV4 "
         "context parallelism is enabled."

python/sglang/srt/layers/moe/mega_moe.py

@@ -24,6 +24,7 @@ import torch
 from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
+from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.model_executor.runner import get_is_capture_mode
@@ -103,7 +104,7 @@ def should_use_mega_moe(moe: DeepseekV2MoE, hidden_states: torch.Tensor) -> bool
         return True
 
     global_num_tokens = get_dp_global_num_tokens()
-    if global_num_tokens:
+    if global_num_tokens and not is_dsa_enable_prefill_cp():
         max_tokens_per_rank = max(global_num_tokens)
     else:
         max_tokens_per_rank = hidden_states.shape[0]

python/sglang/srt/models/deepseek_v4.py

@@ -1783,12 +1783,14 @@ class DeepseekV4DecoderLayer(nn.Module):
             and getattr(self.mlp, "_shared_expert_tp1", False)
         )
         if _use_cp:
-            if get_moe_a2a_backend().is_none():
+            moe_a2a_backend = get_moe_a2a_backend()
+            if moe_a2a_backend.is_none():
                 hidden_states = dsa_cp_gather_hidden_states(hidden_states)
             else:
-                assert get_moe_a2a_backend().is_deepep(), (
-                    "CP requires DeepEP (moe_a2a_backend == deepep). "
-                    "Only DeepEP is tested with CP's per-rank token split."
+                assert moe_a2a_backend.is_deepep() or moe_a2a_backend.is_megamoe(), (
+                    "CP requires DeepEP or megaMoE "
+                    "(moe_a2a_backend == deepep or megamoe). "
+                    f"Got {moe_a2a_backend.value}."
                 )
         elif _use_tp_moe_gather:
             hidden_states, local_hidden_states = (

python/sglang/srt/server_args.py

@@ -4970,9 +4970,13 @@ class ServerArgs:
         elif model_arch in [
             "DeepseekV4ForCausalLM",
         ]:
-            from sglang.srt.arg_groups.deepseek_v4_hook import validate_deepseek_v4_cp
+            from sglang.srt.arg_groups.deepseek_v4_hook import (
+                validate_deepseek_v4_cp,
+                validate_deepseek_v4_mega_moe_token_budget,
+            )
 
             validate_deepseek_v4_cp(self)
+            validate_deepseek_v4_mega_moe_token_budget(self)
 
             # The SM120 marlin fallback moved to the resolution pipeline
             # (arg_groups/overrides.py: _deepseek_v4_sm120_moe), invoked here