[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, value_b = cls._split_num_token_non_padded(
-        value_a = min(tbo_split_token_index, num_token_non_padded)
+            tbo_split_token_index=tbo_split_token_index,
-        value_b = max(0, num_token_non_padded - 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
+                    # Keep idle non-hybrid fabricated rows masked by default.
-                    # treats this rank as empty and starves later layers.
+                    # Hybrid-SSM needs the real count for its state update.
-                    # (num_token_non_padded is None unless moe_ep_size > 1.)
+                    mask_dummy_tokens = (
-                    if self.num_token_non_padded is not None:
+                        not hybrid_ssm and self._original_forward_mode.is_idle()
-                        self.num_token_non_padded.fill_(num_tokens)
+                    )
-                    self.num_token_non_padded_cpu = num_tokens
+                    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]
 

test/registered/unit/batch_overlap/test_tbo_children_dummy_token_mask.py

@@ -0,0 +1,148 @@
+"""Regression: the TBO split must not resurrect masked idle-rank dummy tokens.
+
+Under DP attention an idle rank's batch is rewritten into a fabricated EXTEND
+batch and then padded, and ``prepare_mlp_sync_batch`` masks the fabricated rows
+by setting ``num_token_non_padded(_cpu)`` to 0 so MoE top-k skips them and
+attention returns early. ``TboForwardBatchPreparer`` runs right after that, and
+it used to (a) recompute the children counts from ``len(batch.input_ids)`` --
+the *padded* token count, which restores the dummy rows as real tokens for MoE
+-- and (b) hardcode each child's ``num_token_non_padded_cpu`` to ``None``, so
+the ``real_num_tokens == 0`` attention skip compared ``None == 0`` and never
+fired. Net effect: DeepEP/pplx MAX_LEN + ``--enable-two-batch-overlap``
+silently lost the whole idle-rank optimization.
+
+CPU-only.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_context, get_device, get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _make_extend_batch(*, padded_num_tokens: int, num_token_non_padded_cpu: int):
+    # Only the fields compute_tbo_children_num_token_non_padded reads.
+    return SimpleNamespace(
+        forward_mode=ForwardMode.EXTEND,
+        spec_info=None,
+        tbo_split_seq_index=1,
+        extend_seq_lens_cpu=[4, 4],
+        input_ids=torch.zeros(padded_num_tokens, dtype=torch.long),
+        num_token_non_padded_cpu=num_token_non_padded_cpu,
+    )
+
+
+def _make_decode_capture_batch(*, num_tokens: int):
+    # Decode CUDA-graph capture batches do not populate the CPU mirror.
+    return SimpleNamespace(
+        forward_mode=ForwardMode.DECODE,
+        spec_info=None,
+        tbo_split_seq_index=2,
+        extend_seq_lens_cpu=None,
+        input_ids=torch.zeros(num_tokens, dtype=torch.long),
+        num_token_non_padded_cpu=None,
+    )
+
+
+class TestTboChildrenDummyTokenMask(CustomTestCase):
+    def _children_counts(self, batch):
+        with get_context().override_server_args():
+            with get_device().override(device="cpu"):
+                return (
+                    TboForwardBatchPreparer.compute_tbo_children_num_token_non_padded(
+                        batch
+                    ).tolist()
+                )
+
+    def test_masked_idle_parent_yields_zero_token_children(self):
+        # An idle rank masked to 0 real tokens must split into two 0-token
+        # children even though input_ids still holds the padded dummy rows.
+        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=0)
+        self.assertEqual(self._children_counts(batch), [0, 0])
+
+    def test_padding_is_not_counted_as_real_tokens(self):
+        # A busy rank padded up to a larger bucket must split on its real token
+        # count, not on the padded input_ids length.
+        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=8)
+        self.assertEqual(self._children_counts(batch), [4, 4])
+
+    def test_cpu_pair_matches_device_pair(self):
+        # prepare() derives the children's CPU counts separately from the device
+        # tensor; the two must not drift apart.
+        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=5)
+        cpu_pair = TboForwardBatchPreparer._split_num_token_non_padded(
+            tbo_split_token_index=TboForwardBatchPreparer._compute_split_token_index(
+                batch
+            ),
+            num_token_non_padded=batch.num_token_non_padded_cpu,
+        )
+        self.assertEqual(list(cpu_pair), self._children_counts(batch))
+
+    def test_filter_batch_propagates_cpu_count_to_child(self):
+        # Without this the attention 0-token skip (which reads
+        # num_token_non_padded_cpu) compares None == 0 and never fires.
+        bs = 8
+        parent = ForwardBatch(
+            forward_mode=ForwardMode.TARGET_VERIFY,
+            batch_size=bs,
+            input_ids=torch.zeros(bs, dtype=torch.long),
+            positions=torch.zeros(bs, dtype=torch.long),
+            out_cache_loc=torch.zeros(bs, dtype=torch.long),
+            req_pool_indices=torch.zeros(bs, dtype=torch.long),
+            seq_lens=torch.ones(bs, dtype=torch.int32),
+            seq_lens_cpu=torch.ones(bs, dtype=torch.int32),
+            seq_lens_sum=bs,
+            spec_info=None,
+        )
+        with (
+            get_context().override_server_args(
+                attention_backend="fa3", moe_dense_tp_size=None
+            ),
+            get_parallel().override(attn_tp_size=1),
+        ):
+            child = TboForwardBatchPreparer.filter_batch(
+                parent,
+                start_token_index=0,
+                end_token_index=4,
+                start_seq_index=0,
+                end_seq_index=4,
+                out_num_token_non_padded=torch.tensor(0),
+                out_num_token_non_padded_cpu=0,
+            )
+        self.assertEqual(child.num_token_non_padded_cpu, 0)
+
+    def test_capture_count_falls_back_to_physical_rows(self):
+        batch = _make_decode_capture_batch(num_tokens=8)
+        self.assertEqual(self._children_counts(batch), [2, 6])
+
+    def test_prepare_falls_back_to_physical_rows_for_missing_cpu_count(self):
+        batch = _make_decode_capture_batch(num_tokens=8)
+
+        with (
+            patch.object(
+                TboForwardBatchPreparer,
+                "compute_tbo_children_num_token_non_padded",
+                return_value=torch.tensor([2, 6], dtype=torch.int32),
+            ),
+            patch.object(TboForwardBatchPreparer, "prepare_raw") as prepare_raw,
+        ):
+            TboForwardBatchPreparer.prepare(batch)
+
+        prepare_raw.assert_called_once()
+        self.assertEqual(
+            prepare_raw.call_args.kwargs["tbo_children_num_token_non_padded_cpu"],
+            (2, 6),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_radix_attention.py

@@ -300,6 +300,117 @@ class TestRadixAttentionGraphInterface(CustomTestCase):
         self.assertTrue(torch.all(output[:2] == 3))
         self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
 
+    def test_impl_zero_real_tokens_returns_zeroed_lse(self):
+        # Regression: an idle DP rank whose fabricated EXTEND batch is masked to
+        # 0 real tokens skips attention entirely. The skip must still honor the
+        # LSE return mode -- unified_attention_with_output_and_lse asserts a
+        # tensor comes back, so returning a bare None raised AssertionError as
+        # soon as any 0-real-token call needed LSE (chunked-prefix MHA merge).
+        attention_layer = SimpleNamespace()
+        context = self._new_impl_context([attention_layer], real_num_tokens=0)
+        backend = _RecordingAttentionBackend()
+        query = torch.zeros((4, 2, 3))
+        output = torch.full_like(query, float("nan"))
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            lse = radix_attention_module._unified_attention_with_output_impl(
+                query,
+                query,
+                query,
+                output,
+                False,
+                0,
+                False,
+                True,
+            )
+
+        self.assertEqual(backend.calls, [])
+        self.assertTrue(torch.all(output == 0))
+        # Same shape/dtype the registered fake impl declares, so
+        # unified_attention_with_output_and_lse's `assert lse is not None` holds.
+        self.assertEqual(lse.shape, (4, 2))
+        self.assertEqual(lse.dtype, torch.float32)
+        self.assertTrue(torch.all(lse == 0))
+
+    def test_impl_zero_real_tokens_output_only_returns_none(self):
+        # The 0-token skip must not start returning a tensor on the non-LSE
+        # path: unified_attention_with_output is registered with an inplace
+        # (None-returning) schema.
+        attention_layer = SimpleNamespace()
+        context = self._new_impl_context([attention_layer], real_num_tokens=0)
+        backend = _RecordingAttentionBackend(return_lse=False)
+        query = torch.zeros((4, 2, 3))
+        output = torch.full_like(query, float("nan"))
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            lse = radix_attention_module._unified_attention_with_output_impl(
+                query,
+                query,
+                query,
+                output,
+                False,
+                0,
+                False,
+                False,
+            )
+
+        self.assertIsNone(lse)
+        self.assertEqual(backend.calls, [])
+        self.assertTrue(torch.all(output == 0))
+
+    def test_extra_kwargs_zero_real_tokens_zeroes_output(self):
+        # Regression: attention_with_output_extra_kwargs (Inkling score_mod /
+        # aux_tensors) narrowed to query[:0] and copied output[:0], so with 0
+        # real tokens the preallocated torch.empty output was never written and
+        # its garbage (NaN/Inf) flowed into residuals and MoE routing. Only ROCm
+        # zeroed the padded tail, so every other platform leaked it.
+        attention_layer = SimpleNamespace()
+        context = self._new_impl_context([attention_layer], real_num_tokens=0)
+        backend = _RecordingAttentionBackend(return_lse=False)
+        query = torch.zeros((4, 2, 3))
+        output = torch.full_like(query, float("nan"))
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            radix_attention_module.attention_with_output_extra_kwargs(
+                query,
+                query,
+                query,
+                output,
+                False,
+                0,
+                {},
+            )
+
+        self.assertEqual(backend.calls, [])
+        self.assertTrue(torch.all(output == 0))
+
     def test_lse_fake_impl_declares_shape_and_dtype(self):
         query = torch.empty((5, 3, 7), dtype=torch.float16)
         output = torch.empty_like(query)