diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 30e10abd4..6c3cc6491 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -50,7 +50,7 @@ class ScheduleBatchDisaggregationDecodeMixin: pre_len = len(req.prefix_indices) chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ - : req.extend_input_len + pre_len : pre_len + req.extend_input_len ] assert ( offset + req.extend_input_len <= total_size @@ -96,7 +96,7 @@ class ScheduleBatchDisaggregationDecodeMixin: seq_len, int(getattr(req, "cached_tokens", 0) or 0), req.req_pool_idx, - 0, + pre_len, pre_len, int(chunk.numel()), ) diff --git a/test/registered/unit/disaggregation/test_decode_queue_compaction.py b/test/registered/unit/disaggregation/test_decode_queue_compaction.py index e78d02e2a..7a8bbb30d 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_compaction.py +++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py @@ -14,6 +14,9 @@ from sglang.srt.disaggregation.decode import ( DecodeTransferQueue, SchedulerDisaggregationDecodeMixin, ) +from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( + ScheduleBatchDisaggregationDecodeMixin, +) from sglang.srt.disaggregation.utils import ( DisaggregationMode, ReqToMetadataIdxAllocator, @@ -728,6 +731,52 @@ class TestDecodeQueueCompaction(CustomTestCase): for req in captured["reqs"]: self.assertGreaterEqual(req.metadata_buffer_index, 20) + def test_prepare_for_prebuilt_uses_suffix_cache_locs_after_cache_hit(self): + """Decode prebuilt must write new-token locs, not prefix locs. + + Prefill transfers the full prompt KV to decode. During the first decode + prebuilt step, `prefix_indices` covers the cached prompt prefix and + `extend_input_len` covers only the prompt suffix that still needs a + local decode forward. The output cache loc tensor must therefore slice + req_to_token at [pre_len : pre_len + extend_input_len]. + """ + + req = FakeReq("cache-hit", 11) + req.req_pool_idx = 0 + req.prefix_indices = torch.arange(100, 104, dtype=torch.int64) + req.extend_input_len = 3 + req.fill_ids = [10, 11, 12, 13, 14, 15, 16] + req.origin_input_ids = list(req.fill_ids) + req.output_ids = [] + req.retracted_stain = False + req.already_computed = len(req.prefix_indices) + req.top_logprobs_num = 0 + req.token_ids_logprob = None + + batch = cast(Any, SimpleNamespace()) + batch.reqs = [req] + batch.req_to_token_pool = SimpleNamespace( + req_to_token=torch.tensor( + [[1000, 1001, 1002, 1003, 2000, 2001, 2002]], + dtype=torch.int64, + ) + ) + batch.device = "cpu" + batch.return_logprob = False + batch.model_config = SimpleNamespace(vocab_size=32000) + batch.multimodal_inputs = None + + with patch( + "sglang.srt.disaggregation.decode_schedule_batch_mixin." + "SamplingBatchInfo.from_schedule_batch", + return_value=SimpleNamespace(), + ): + ScheduleBatchDisaggregationDecodeMixin.prepare_for_prebuilt(batch) + + self.assertEqual(batch.prefix_lens, [4]) + self.assertEqual(batch.extend_lens, [3]) + self.assertEqual(batch.out_cache_loc.tolist(), [2000, 2001, 2002]) + def test_get_new_prebuilt_batch_frees_metadata_on_prebuilt_error(self): allocator = FakeAllocator() scheduler = cast(Any, SimpleNamespace())