[Performance] Reduce idle DP work in breakable prefill CUDA graphs (#33871)

python/sglang/srt/model_executor/forward_batch_info.py

@@ -304,6 +304,17 @@ def compute_local_num_token_non_padded_cpu(
     return min(max(global_num_token_non_padded - rank_offset, 0), tokens_per_rank)
 
 
+def prefill_graph_tolerates_sum_len() -> bool:
+    """Whether MegaMoE may replay prefill graphs with local shapes."""
+    from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+    from sglang.srt.layers.moe.utils import get_moe_a2a_backend
+    from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
+
+    if not get_moe_a2a_backend().is_megamoe():
+        return False
+    return not (is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled())
+
+
 @dataclass
 class DSV4OutCacheLoc:
     """Per-forward-pass KV cache allocation for DeepSeek-V4 on NPU.
@@ -1329,6 +1340,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             and self.is_extend_in_batch
             and prefill_cg.bs
             and max(global_num_tokens) <= max(prefill_cg.bs)
+            and not prefill_graph_tolerates_sum_len()
         ):
             dp_padding_mode = DpPaddingMode.MAX_LEN
         self.dp_padding_mode = dp_padding_mode
@@ -1431,12 +1443,19 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                     self.extend_seq_lens_cpu = [int(num_tokens)]
                     self.extend_logprob_start_lens_cpu = [0]
                     bs = self.batch_size = 1
-                    # Count the dummy tokens as real, else MoE topk/all-to-all
-                    # treats this rank as empty and starves later layers.
-                    # (num_token_non_padded is None unless moe_ep_size > 1.)
-                    if self.num_token_non_padded is not None:
-                        self.num_token_non_padded.fill_(num_tokens)
-                    self.num_token_non_padded_cpu = num_tokens
+                    # Keep idle non-hybrid fabricated rows masked by default.
+                    # Hybrid-SSM needs the real count for its state update.
+                    mask_dummy_tokens = (
+                        not hybrid_ssm and self._original_forward_mode.is_idle()
+                    )
+                    if mask_dummy_tokens:
+                        if self.num_token_non_padded is not None:
+                            self.num_token_non_padded.fill_(0)
+                        self.num_token_non_padded_cpu = 0
+                    else:
+                        if self.num_token_non_padded is not None:
+                            self.num_token_non_padded.fill_(num_tokens)
+                        self.num_token_non_padded_cpu = num_tokens
                 else:
                     self.extend_num_tokens = bs
                     self.extend_seq_lens = torch.full_like(self.seq_lens, 1)

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -84,6 +84,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     PPProxyTensors,
     compute_local_num_token_non_padded,
     enable_num_token_non_padded,
+    prefill_graph_tolerates_sum_len,
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
 from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
@@ -777,9 +778,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # DSV4 DP attention / DeepEP collectives need every DP rank to enter
         # the same replay path. Sparse-DP batches (one or more ranks with
         # zero local tokens) fall back to eager to avoid hanging ranks.
+        # MegaMoE is exempt (prefill_graph_tolerates_sum_len): its idle ranks
+        # still execute MegaMoE with 0 tokens, so per-rank SUM_LEN buckets stay
+        # collective-safe and need no eager fallback.
         global_num_tokens = forward_batch.global_num_tokens_cpu
         if global_num_tokens is None:
             return False
+        if prefill_graph_tolerates_sum_len():
+            return False
         return len(global_num_tokens) > 1 and any(
             int(num_tokens) == 0 for num_tokens in global_num_tokens
         )