From 718bd39fe08b42e0095daa571510e47c40c588d2 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 3 Sep 2026 07:14:16 +0800 Subject: [PATCH] [Fix] DP attention: correct the decode->extend prefix off-by-one (#37505) --- python/sglang/srt/managers/schedule_batch.py | 43 ++++--- .../managers/scheduler_components/dp_attn.py | 4 + .../scheduler_components/test_dp_attn.py | 55 ++++++++ ...schedule_batch_convert_decode_to_extend.py | 119 ++++++++++++++++++ .../test_schedule_batch_out_of_place.py | 4 +- 5 files changed, 202 insertions(+), 23 deletions(-) create mode 100644 test/registered/unit/managers/test_schedule_batch_convert_decode_to_extend.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 2323aa7af..f586bc277 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2886,10 +2886,18 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.forward_mode = ForwardMode.MIXED running_bs = running_batch.batch_size() - for req in running_batch.reqs: + # Same invariant as convert_decode_to_extend: the caller ran + # prepare_for_decode, so a tail's prefix is its row length - 1. + if self.spec_algorithm.is_none(): + running_prefix_lens = [s - 1 for s in running_batch.seq_lens_cpu.tolist()] + else: + # Spec rows sit at the committed base; seq_lens is rebuilt below. + running_prefix_lens = [r.seqlen - 1 for r in running_batch.reqs] + for req, prefix_len in zip( + running_batch.reqs, running_prefix_lens, strict=True + ): req._refresh_fill_ids() - full_len = len(req.full_untruncated_fill_ids) - req.set_extend_range(full_len - 1, full_len) + req.set_extend_range(prefix_len, prefix_len + 1) # Decode tokens of the running portion live in future_map.output_tokens_buf. self.input_ids = None @@ -2937,18 +2945,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): merged[-running_bs:] = tail_base + 1 self.seq_lens = merged - # For overlap scheduler, the output_ids has one step delay; - # spec tail request state carries no delay in either mode. - if self.spec_algorithm.is_none(): - delta = 0 if self.enable_overlap else -1 - else: - delta = -1 - # NOTE: prefix_indices is what has been cached, but we don't cache each decode step - self.prefix_lens = self.prefix_lens + [ - len(r.origin_input_ids) + len(r.output_ids) + delta - for r in running_batch.reqs - ] + self.prefix_lens = self.prefix_lens + running_prefix_lens self.extend_lens = self.extend_lens + [1] * running_bs self.extend_num_tokens = self.extend_num_tokens + running_bs # TODO (lianmin): Revisit this. It should be seq_len - 1 @@ -2971,16 +2969,17 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Also stale residue; None keeps the prefill result path from # re-reporting old prefill stats for what is decode work. self.prefill_stats = None - for req in self.reqs: + # A 1-token extend's position is arange(prefix, prefix + 1), so prefix + # must be seq_len - 1; output_ids trails the row by one or zero and + # cannot stand in for it. Rows past bs are a beam tail, not requests. + seq_lens = self.seq_lens_cpu[:bs].tolist() + for req, seq_len in zip(self.reqs, seq_lens, strict=True): req._refresh_fill_ids() - full_len = len(req.full_untruncated_fill_ids) - req.set_extend_range(full_len - 1, full_len) + # end runs one past full_untruncated_fill_ids while output_ids + # trails; safe only while decoding_reqs suppresses cache_unfinished_req. + req.set_extend_range(seq_len - 1, seq_len) - # Same one-step output_ids delay handling as mix_with_running. - delta = 0 if self.enable_overlap else -1 - self.prefix_lens = [ - len(r.origin_input_ids) + len(r.output_ids) + delta for r in self.reqs - ] + self.prefix_lens = [seq_len - 1 for seq_len in seq_lens] self.extend_lens = [1] * bs self.extend_num_tokens = bs self.extend_logprob_start_lens = [0] * bs diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index bffbfab0f..728e18814 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -309,6 +309,10 @@ def _local_prefill_cuda_graph_vote( and not local_batch.return_logprob # Grammar FSMs advance through the decode result path only. and not local_batch.has_grammar + # A converted batch takes the prefill result path, which commits beam + # requests per-req rather than through the batch decode fold; member + # rows also have no req of their own for the reqs-aligned extend lists. + and all(r.beam_group is None for r in local_batch.reqs) # Small-bucket BCG replays amplify the a2a EP logits drift (#30898) # into an accuracy loss. and get_moe_a2a_backend().is_none() diff --git a/test/registered/unit/managers/scheduler_components/test_dp_attn.py b/test/registered/unit/managers/scheduler_components/test_dp_attn.py index 59e792266..f69979861 100644 --- a/test/registered/unit/managers/scheduler_components/test_dp_attn.py +++ b/test/registered/unit/managers/scheduler_components/test_dp_attn.py @@ -75,5 +75,60 @@ class TestDPAttnSchedulerMetadata(CustomTestCase): ) +class TestDecodeToExtendConversionVote(CustomTestCase): + """A decode batch votes for the prefill graph only when its 1-token extend + view can represent every row. Beam requests cannot: the converted batch + takes the prefill result path, which commits them per-req instead of + through the batch decode fold, and member rows carry no req.""" + + def _vote(self, *, beam): + runner = Mock(spec=dp_attn.PrefillCudaGraphRunner) + runner.enable_lora = False + runner.can_replay_locally.return_value = True + batch = SimpleNamespace( + forward_mode=ForwardMode.DECODE, + batch_size=lambda: 2, + return_logprob=False, + has_grammar=False, + reqs=[ + SimpleNamespace(beam_group=Mock() if beam else None), + SimpleNamespace(beam_group=None), + ], + ) + with ( + patch.object( + dp_attn, "get_moe_a2a_backend", return_value=Mock(is_none=lambda: True) + ), + patch.object(dp_attn, "uses_ssm_state", return_value=False), + patch.object( + dp_attn, + "get_memory", + return_value=SimpleNamespace(enable_hisparse=False), + ), + patch.object( + dp_attn, + "get_exec", + return_value=SimpleNamespace( + overlap=SimpleNamespace(enable_two_batch_overlap=False) + ), + ), + patch.object(dp_attn, "get_cp_strategy", return_value=None), + ): + return dp_attn._local_prefill_cuda_graph_vote( + local_batch=batch, + prefill_graph_runner=runner, + coordinated_prefill=True, + breakable_prefill=True, + spec_algorithm=SpeculativeAlgorithm.NONE, + model_config=object(), + ) + + def test_plain_decode_batch_votes_for_conversion(self): + self.assertTrue(self._vote(beam=False)) + + def test_beam_request_blocks_conversion(self): + self.assertFalse(self._vote(beam=True)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_schedule_batch_convert_decode_to_extend.py b/test/registered/unit/managers/test_schedule_batch_convert_decode_to_extend.py new file mode 100644 index 000000000..929e8a4fc --- /dev/null +++ b/test/registered/unit/managers/test_schedule_batch_convert_decode_to_extend.py @@ -0,0 +1,119 @@ +"""The dp decode->extend view must place the new token at seq_len - 1. + +Deriving the prefix from len(origin_input_ids) + len(output_ids) put RoPE one +position past the row's own KV slot whenever the overlap output_ids lag was +drained. +""" + +import unittest +from array import array + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.schedule_batch import ( # noqa: E402 + ForwardMode, + Req, + ScheduleBatch, +) + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _FakeReq: + """Carries the fill-id state convert_decode_to_extend touches, with the + real Req methods so the array bookkeeping is not re-implemented here.""" + + _refresh_fill_ids = Req._refresh_fill_ids + set_extend_range = Req.set_extend_range + + def __init__(self, *, num_prompt_tokens: int, num_output_tokens: int): + self.origin_input_ids = array("l", range(num_prompt_tokens)) + self.output_ids = list(range(num_output_tokens)) + self.full_untruncated_fill_ids = array("l", self.origin_input_ids) + self.extend_range = None + self.beam_group = None + + +def _make_converted_batch(*, rows, output_lag, extra_rows=0): + """rows: (num_prompt_tokens, seq_len) per request. output_lag: how many + tokens output_ids trails seq_len by (1 in steady decode, 0 once drained).""" + reqs = [ + _FakeReq( + num_prompt_tokens=num_prompt_tokens, + num_output_tokens=seq_len - num_prompt_tokens - output_lag, + ) + for num_prompt_tokens, seq_len in rows + ] + seq_lens = [seq_len for _, seq_len in rows] + batch = ScheduleBatch(reqs=reqs) + batch.forward_mode = ForwardMode.DECODE + batch.enable_overlap = True + # A beam tail appends rows after the reqs-aligned ones; mimic it by + # repeating the last row, which is what append_beam_tail does. + tail = [seq_lens[-1]] * extra_rows + batch.seq_lens_cpu = torch.tensor(seq_lens + tail, dtype=torch.int64) + batch.convert_decode_to_extend() + return batch, seq_lens + + +class TestConvertDecodeToExtendGeometry(CustomTestCase): + def _assert_geometry(self, batch, seq_lens): + self.assertEqual(batch.forward_mode, ForwardMode.EXTEND) + self.assertEqual(batch.prefix_lens, [s - 1 for s in seq_lens]) + self.assertEqual(batch.extend_lens, [1] * len(seq_lens)) + self.assertEqual(batch.extend_num_tokens, len(seq_lens)) + for req, seq_len in zip(batch.reqs, seq_lens, strict=True): + self.assertEqual(tuple(req.extend_range), (seq_len - 1, seq_len)) + # What the attention path actually consumes: arange(prefix, prefix+len) + # must land on the slot prepare_for_decode allocated, at seq_len - 1. + for prefix_len, extend_len, seq_len in zip( + batch.prefix_lens, batch.extend_lens, seq_lens, strict=True + ): + self.assertEqual(prefix_len + extend_len, seq_len) + + def test_geometry_with_the_overlap_lag_present(self): + """Steady decode: the previous step's token is not in output_ids yet.""" + batch, seq_lens = _make_converted_batch( + rows=[(6, 8), (96, 97), (142, 143)], output_lag=1 + ) + self._assert_geometry(batch, seq_lens) + + def test_geometry_once_the_overlap_lag_is_drained(self): + """An iteration that ran prefill instead of decode lets the pending + result land, so origin + output_ids reaches seq_len. The pre-fix + formula returned prefix_len == seq_len here, i.e. RoPE one past the + row's own KV slot.""" + batch, seq_lens = _make_converted_batch( + rows=[(6, 8), (96, 97), (142, 143)], output_lag=0 + ) + self._assert_geometry(batch, seq_lens) + + def test_prefix_lens_stays_reqs_aligned_under_a_beam_tail(self): + """seq_lens_cpu carries beam member rows with no req of their own; the + reqs-aligned lists must not grow to the row count.""" + batch, seq_lens = _make_converted_batch( + rows=[(6, 8), (96, 97)], output_lag=1, extra_rows=3 + ) + self.assertEqual(len(batch.seq_lens_cpu), len(seq_lens) + 3) + self.assertEqual(len(batch.prefix_lens), len(batch.reqs)) + self._assert_geometry(batch, seq_lens) + + def test_short_seq_lens_fails_loudly(self): + """A row/req divergence the slice cannot explain must raise, not + silently truncate the way a plain zip would.""" + reqs = [_FakeReq(num_prompt_tokens=6, num_output_tokens=1) for _ in range(3)] + batch = ScheduleBatch(reqs=reqs) + batch.forward_mode = ForwardMode.DECODE + batch.enable_overlap = True + batch.seq_lens_cpu = torch.tensor([8, 8], dtype=torch.int64) + with self.assertRaises(ValueError): + batch.convert_decode_to_extend() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_schedule_batch_out_of_place.py b/test/registered/unit/managers/test_schedule_batch_out_of_place.py index 4fc0c562e..c91500ff4 100644 --- a/test/registered/unit/managers/test_schedule_batch_out_of_place.py +++ b/test/registered/unit/managers/test_schedule_batch_out_of_place.py @@ -145,6 +145,8 @@ class TestMixWithRunningOutOfPlace(unittest.TestCase): return_logprob=False, forward_mode=ForwardMode.DECODE, out_cache_loc=torch.arange(6, 7, dtype=torch.int64), + # prepare_for_decode ran: the row already counts this step's token. + seq_lens_cpu=torch.tensor([6], dtype=torch.int64), ) extend_prefix_before = extend_batch.prefix_lens @@ -160,7 +162,7 @@ class TestMixWithRunningOutOfPlace(unittest.TestCase): self.assertTrue( torch.equal(extend_batch.out_cache_loc, torch.arange(7, dtype=torch.int64)) ) - # delta is -1 without overlap: 4 origin + 2 output - 1 + # The decode tail's prefix is its row length minus this step's token. self.assertEqual(extend_batch.prefix_lens, [0, 0, 5]) self.assertEqual(extend_batch.extend_lens, [3, 3, 1]) self.assertEqual(extend_batch.extend_num_tokens, 7)