minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
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)