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

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -513,12 +513,24 @@ class TboForwardBatchPreparer:
             cls.compute_tbo_children_num_token_non_padded(batch)
         )
         cls.prepare_raw(
-            batch, tbo_children_num_token_non_padded=tbo_children_num_token_non_padded
+            batch,
+            tbo_children_num_token_non_padded=tbo_children_num_token_non_padded,
+            # Eager split: the children can carry a CPU count too, so the
+            # attention 0-token skip (which reads num_token_non_padded_cpu)
+            # survives the split. The cuda-graph plugin path below leaves this
+            # None because its device buffer is refreshed per replay.
+            tbo_children_num_token_non_padded_cpu=cls._split_num_token_non_padded(
+                tbo_split_token_index=cls._compute_split_token_index(batch),
+                num_token_non_padded=cls._get_num_token_non_padded_cpu(batch),
+            ),
         )
 
     @classmethod
     def prepare_raw(
-        cls, batch: ForwardBatch, tbo_children_num_token_non_padded: torch.Tensor
+        cls,
+        batch: ForwardBatch,
+        tbo_children_num_token_non_padded: torch.Tensor,
+        tbo_children_num_token_non_padded_cpu: Optional[tuple[int, int]] = None,
     ):
         from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
 
@@ -548,6 +560,9 @@ class TboForwardBatchPreparer:
         [out_num_token_non_padded_a, out_num_token_non_padded_b] = (
             tbo_children_num_token_non_padded
         )
+        out_num_token_non_padded_cpu_a, out_num_token_non_padded_cpu_b = (
+            tbo_children_num_token_non_padded_cpu or (None, None)
+        )
 
         child_a = cls.filter_batch(
             batch,
@@ -560,6 +575,7 @@ class TboForwardBatchPreparer:
                 else batch.tbo_split_seq_index
             ),
             out_num_token_non_padded=out_num_token_non_padded_a,
+            out_num_token_non_padded_cpu=out_num_token_non_padded_cpu_a,
         )
         child_b = cls.filter_batch(
             batch,
@@ -568,6 +584,7 @@ class TboForwardBatchPreparer:
             start_seq_index=batch.tbo_split_seq_index,
             end_seq_index=batch.batch_size,
             out_num_token_non_padded=out_num_token_non_padded_b,
+            out_num_token_non_padded_cpu=out_num_token_non_padded_cpu_b,
         )
 
         if is_enable_two_chunk:
@@ -655,6 +672,7 @@ class TboForwardBatchPreparer:
         start_seq_index: int,
         end_seq_index: int,
         out_num_token_non_padded: torch.Tensor,
+        out_num_token_non_padded_cpu: Optional[int] = None,
     ):
         assert (
             end_token_index >= start_token_index
@@ -788,7 +806,7 @@ class TboForwardBatchPreparer:
                 extend_num_tokens=extend_num_tokens,
                 num_token_non_padded=out_num_token_non_padded,
                 # TODO: handle it when we need TBO + DeepSeek V3.2
-                num_token_non_padded_cpu=None,
+                num_token_non_padded_cpu=out_num_token_non_padded_cpu,
                 tbo_split_seq_index=None,
                 tbo_parent_token_range=(start_token_index, end_token_index),
                 tbo_children=None,
@@ -835,20 +853,43 @@ class TboForwardBatchPreparer:
     def compute_tbo_children_num_token_non_padded(cls, batch: ForwardBatch):
         return cls.compute_tbo_children_num_token_non_padded_raw(
             tbo_split_token_index=cls._compute_split_token_index(batch),
-            num_token_non_padded=len(batch.input_ids),
+            # Prefer the parent CPU count: len(input_ids) is the padded
+            # (MAX_LEN) count and would undo the idle-rank dummy-token mask.
+            # The resolver falls back to physical rows only for capture
+            # batches that intentionally leave the CPU mirror unset.
+            num_token_non_padded=cls._get_num_token_non_padded_cpu(batch),
         )
 
+    @staticmethod
+    def _get_num_token_non_padded_cpu(batch: ForwardBatch) -> int:
+        num_token_non_padded = (
+            batch.num_token_non_padded_cpu
+            if batch.num_token_non_padded_cpu is not None
+            else len(batch.input_ids)
+        )
+        return num_token_non_padded
+
     @classmethod
     def compute_tbo_children_num_token_non_padded_raw(
         cls, tbo_split_token_index: int, num_token_non_padded: int
     ):
-        # TODO we may make padding on both sub-batches to make it slightly more balanced
-        value_a = min(tbo_split_token_index, num_token_non_padded)
-        value_b = max(0, num_token_non_padded - tbo_split_token_index)
+        value_a, value_b = cls._split_num_token_non_padded(
+            tbo_split_token_index=tbo_split_token_index,
+            num_token_non_padded=num_token_non_padded,
+        )
         return torch.tensor([value_a, value_b], dtype=torch.int32).to(
             device=get_device().device, non_blocking=True
         )
 
+    @staticmethod
+    def _split_num_token_non_padded(
+        *, tbo_split_token_index: int, num_token_non_padded: int
+    ) -> tuple[int, int]:
+        # TODO we may make padding on both sub-batches to make it slightly more balanced
+        value_a = min(tbo_split_token_index, num_token_non_padded)
+        value_b = max(0, num_token_non_padded - tbo_split_token_index)
+        return value_a, value_b
+
     @classmethod
     def _compute_split_token_index(cls, batch: ForwardBatch):
         token_num_per_seq = get_token_num_per_seq(

python/sglang/srt/layers/radix_attention.py

@@ -69,6 +69,13 @@ def _zero_padded_pcg_tail(buf: torch.Tensor, context) -> None:
         buf.view(first_dim, elems_per_token)[actual_tokens:].zero_()
 
 
+def _zero_skipped_attn_outputs(*bufs: Optional[torch.Tensor]) -> None:
+    """Zero outputs when an idle DP rank skips attention work."""
+    for buf in bufs:
+        if buf is not None:
+            buf.zero_()
+
+
 if TYPE_CHECKING:
     from sglang.srt.layers.quantization.base_config import QuantizationConfig
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -319,6 +326,18 @@ def _unified_attention_with_output_impl(
     if key_value_num_tokens is None:
         key_value_num_tokens = real_query_num_tokens
 
+    if real_query_num_tokens == 0:
+        _zero_skipped_attn_outputs(output)
+        if return_lse:
+            # unified_attention_with_output_and_lse asserts a tensor comes back.
+            # Match _unified_attention_with_output_and_lse_fake's meta shape and
+            # the padded LSE the normal path returns below (padded row count,
+            # i.e. query before narrowing).
+            return query.new_zeros(
+                (query.shape[0], query.shape[1]), dtype=torch.float32
+            )
+        return None
+
     query = query[:real_query_num_tokens]
     if key is not None:
         key = key[:key_value_num_tokens]
@@ -510,6 +529,10 @@ def unified_sparse_attention_with_output(
     attention_layer = context.attention_layers[layer_id]
     real_num_tokens = forward_batch.num_token_non_padded_cpu
 
+    if real_num_tokens == 0:
+        _zero_skipped_attn_outputs(attn_out, idx_out)
+        return
+
     query = query[:real_num_tokens]
     if key is not None:
         key = key[:real_num_tokens]
@@ -577,6 +600,10 @@ def attention_with_output_extra_kwargs(
     attention_layer = context.attention_layers[layer_id]
     real_num_tokens = forward_batch.num_token_non_padded_cpu
 
+    if real_num_tokens == 0:
+        _zero_skipped_attn_outputs(output)
+        return
+
     query = query[:real_num_tokens]
     if key is not None:
         key = key[:real_num_tokens]

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
         )

python/sglang/srt/models/deepseek_v4.py

@@ -530,6 +530,10 @@ def deepseek_v4_attention_with_output(
     attention_layer = attention_layers[layer_id]
     real_num_tokens = forward_batch.num_token_non_padded_cpu
 
+    if real_num_tokens == 0:
+        output.zero_()
+        return
+
     query = query[:real_num_tokens]
     key_value = key_value[:real_num_tokens]