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

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)