A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 22:15 UTC
[Core] Reuse created spec tokens lists to mitigate GC cost (#28917)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Nov 19, 2025 at 11:20 UTC
3319a493fcc3e4733382f0dc812184234e9c3dcb
2 files changed
+14
-7
vllm/v1/worker/gpu_input_batch.py
+12
-6
@@ -251,7 +251,7 @@ class InputBatch:
251
self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids
252
253
# Store last speculative tokens for sampler.
254
- self.spec_token_ids: list[list[int] | None] = []
254
+ self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)]
255
256
# This is updated each time the batch constituents change.
257
self.sampling_metadata = self._make_sampling_metadata()
@@ -313,7 +313,7 @@ class InputBatch:
313
else:
314
self._req_ids[req_index] = req_id
315
self.req_output_token_ids[req_index] = request.output_token_ids
316
- self.spec_token_ids[req_index] = []
316
+ self.spec_token_ids[req_index].clear()
317
318
self.req_id_to_index[req_id] = req_index
319
@@ -462,7 +462,7 @@ class InputBatch:
462
self.batch_update_builder.removed_append(req_index)
463
self._req_ids[req_index] = None
464
self.req_output_token_ids[req_index] = None
465
- self.spec_token_ids[req_index] = None
465
+ self.spec_token_ids[req_index].clear()
466
467
# LoRA
468
lora_id = self.request_lora_mapping[req_index]
@@ -654,9 +654,15 @@ class InputBatch:
654
self.req_output_token_ids[last_req_index] = None
655
self.req_id_to_index[req_id] = empty_index
656
657
- spec_token_ids = self.spec_token_ids[last_req_index]
658
- self.spec_token_ids[empty_index] = spec_token_ids
659
- self.spec_token_ids[last_req_index] = None
657
+ if last_req_index != empty_index:
658
+ (
659
+ self.spec_token_ids[last_req_index],
660
+ self.spec_token_ids[empty_index],
661
+ ) = (
662
+ self.spec_token_ids[empty_index],
663
+ self.spec_token_ids[last_req_index],
664
+ )
665
+ self.spec_token_ids[last_req_index].clear()
666
667
num_tokens = self.num_tokens[last_req_index]
668
self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
vllm/v1/worker/gpu_model_runner.py
+2
-1
@@ -892,7 +892,8 @@ class GPUModelRunner(
892
# conform to the schema. This can result in
893
# scheduler_output.scheduled_spec_decode_tokens being empty,
894
# even when speculative decoding is enabled.
895
- self.input_batch.spec_token_ids[req_index] = spec_token_ids
895
+ self.input_batch.spec_token_ids[req_index].clear()
896
+ self.input_batch.spec_token_ids[req_index].extend(spec_token_ids)
897
898
# there are no draft tokens with async scheduling,
899
# we clear the spec_decoding info in scheduler_output and