diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 55ed74897..a87f60677 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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( diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 3ec27dc72..fe2178cdc 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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] diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 85e8157e0..b2515aaaa 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index b10ca8f09..12e12c6eb 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 ) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 30ee337f9..8052e150d 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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] diff --git a/test/registered/unit/batch_overlap/test_tbo_children_dummy_token_mask.py b/test/registered/unit/batch_overlap/test_tbo_children_dummy_token_mask.py new file mode 100644 index 000000000..2a767b86a --- /dev/null +++ b/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() diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index aec9f3254..106f1839e 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/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)