[PD] Avoid unused PREBUILT prompt tensor transfer (#35070)

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]
+        # PREBUILT never enters a model forward. Keep the legacy scalar metadata,
-        extend_num_tokens = sum(len(ids) for ids in input_ids)
+        # 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(
+        # The first decode input and speculative extras are seeded by
-            sum(input_ids, array("q")), dtype=torch.int32, device=self.device
+        # 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
         )

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)