[Fix] DP attention: correct the decode->extend prefix off-by-one (#37505)

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

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()

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()

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()

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)