diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 9eaedb980..6488e2eea 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -from array import array from http import HTTPStatus from typing import TYPE_CHECKING, List @@ -30,14 +29,16 @@ class ScheduleBatchDisaggregationDecodeMixin: self.forward_mode = ForwardMode.PREBUILT reqs = self.reqs - input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs] - extend_num_tokens = sum(len(ids) for ids in input_ids) + # PREBUILT never enters a model forward. Keep the legacy scalar metadata, + # but do not flatten and copy every transferred prompt to the GPU only to + # discard it before the first decode step. seq_lens = [] pre_lens = [] req_pool_indices = [] # Pre-calculate total size total_size = sum(req.extend_range.length for req in reqs) + extend_num_tokens = total_size out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device) # Fill the tensor in one pass @@ -77,9 +78,10 @@ class ScheduleBatchDisaggregationDecodeMixin: pre_lens.append(pre_len) # Set fields - self.input_ids = torch.tensor( - sum(input_ids, array("q")), dtype=torch.int32, device=self.device - ) + # The first decode input and speculative extras are seeded by + # process_prebuilt through FutureMap, so merge/forward-entry rebuilds + # input_ids from the relay. + self.input_ids = None self.req_pool_indices = torch.tensor( req_pool_indices, dtype=torch.int64, device=self.device ) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 483dcdaf6..11bb34013 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -18,6 +18,9 @@ from sglang.srt.disaggregation.common.utils import ( unpack_int_lists, unpack_list_of_buffers, ) +from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( + ScheduleBatchDisaggregationDecodeMixin, +) from sglang.srt.disaggregation.mooncake.conn import ( KVArgsRegisterInfo, MooncakeKVManager, @@ -95,6 +98,40 @@ class TestDisaggregationWire(unittest.TestCase): packed = pack_int_lists([[]], "I") self.assertEqual(unpack_int_lists(packed, "I"), [[]]) + def test_prebuilt_skips_unused_prompt_tensor(self): + req = SimpleNamespace( + req_pool_idx=0, + prefix_indices=[0, 1], + extend_range=SimpleNamespace(length=3), + origin_input_ids=[0, 1, 2, 3, 4], + output_ids=[], + retracted_stain=True, + is_retracted=True, + multimodal_inputs=None, + get_fill_ids=Mock(side_effect=AssertionError("prompt should not be read")), + ) + batch = SimpleNamespace( + reqs=[req], + device="cpu", + req_to_token_pool=SimpleNamespace( + req_to_token=torch.arange(5, dtype=torch.int64).reshape(1, 5) + ), + return_logprob=False, + model_config=SimpleNamespace(vocab_size=32), + ) + + with patch( + "sglang.srt.disaggregation.decode_schedule_batch_mixin." + "SamplingBatchInfo.from_schedule_batch", + return_value=Mock(), + ): + ScheduleBatchDisaggregationDecodeMixin.prepare_for_prebuilt(batch) + + self.assertIsNone(batch.input_ids) + self.assertEqual(batch.extend_num_tokens, 3) + self.assertTrue(torch.equal(batch.out_cache_loc, torch.tensor([2, 3, 4]))) + req.get_fill_ids.assert_not_called() + def test_list_of_buffers_roundtrip(self): bufs = [b"abc", b"", b"de", b"x" * 17] self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs)