minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Promisewjx
2026-08-27 10:09:54 +08:00
committed by GitHub
parent 4f59a8dcfa
commit 9d07b9e227
7 changed files with 369 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/two_batch_overlap.py
+48 -7
View File
@@ -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
+27
View File
@@ -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
+25 -6
View File
@@ -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
+6
View File
@@ -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
+4
View File
@@ -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
+148
View File
@@ -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
+111
View File
@@ -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)