From 3348f486892915045816ddf71bdaf351f8ab8d3a Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 17 Nov 2025 23:10:37 -0800 Subject: [PATCH 1/3] [Core] Reuse created spec tokens lists to mitigate GC cost Signed-off-by: Jialin Ouyang --- vllm/v1/worker/gpu_input_batch.py | 8 ++++---- vllm/v1/worker/gpu_model_runner.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 7cf6afa3fc37..34b0fd9e9012 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -251,7 +251,7 @@ def __init__( self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids # Store last speculative tokens for sampler. - self.spec_token_ids: list[list[int] | None] = [] + self.spec_token_ids: list[list[int]] = [] # This is updated each time the batch constituents change. self.sampling_metadata = self._make_sampling_metadata() @@ -313,7 +313,7 @@ def add_request( else: self._req_ids[req_index] = req_id self.req_output_token_ids[req_index] = request.output_token_ids - self.spec_token_ids[req_index] = [] + self.spec_token_ids[req_index].clear() self.req_id_to_index[req_id] = req_index @@ -462,7 +462,7 @@ def remove_request(self, req_id: str) -> int | None: self.batch_update_builder.removed_append(req_index) self._req_ids[req_index] = None self.req_output_token_ids[req_index] = None - self.spec_token_ids[req_index] = None + self.spec_token_ids[req_index].clear() # LoRA lora_id = self.request_lora_mapping[req_index] @@ -656,7 +656,7 @@ def condense(self) -> None: spec_token_ids = self.spec_token_ids[last_req_index] self.spec_token_ids[empty_index] = spec_token_ids - self.spec_token_ids[last_req_index] = None + self.spec_token_ids[last_req_index].clear() num_tokens = self.num_tokens[last_req_index] self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0102ca4739ad..f74f3732b76f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -891,7 +891,8 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None: # conform to the schema. This can result in # scheduler_output.scheduled_spec_decode_tokens being empty, # even when speculative decoding is enabled. - self.input_batch.spec_token_ids[req_index] = spec_token_ids + self.input_batch.spec_token_ids[req_index].clear() + self.input_batch.spec_token_ids[req_index].extend(spec_token_ids) # there are no draft tokens with async scheduling, # we clear the spec_decoding info in scheduler_output and From d0c915c700d391f19fc4522d06f9a1f7303bdb3b Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Mon, 17 Nov 2025 23:17:24 -0800 Subject: [PATCH 2/3] Fix errors Signed-off-by: Jialin Ouyang --- vllm/v1/worker/gpu_input_batch.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 34b0fd9e9012..113e0e0c5a38 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -654,9 +654,15 @@ def condense(self) -> None: self.req_output_token_ids[last_req_index] = None self.req_id_to_index[req_id] = empty_index - spec_token_ids = self.spec_token_ids[last_req_index] - self.spec_token_ids[empty_index] = spec_token_ids - self.spec_token_ids[last_req_index].clear() + if last_req_index != empty_index: + ( + self.spec_token_ids[last_req_index], + self.spec_token_ids[empty_index], + ) = ( + self.spec_token_ids[empty_index], + self.spec_token_ids[last_req_index], + ) + self.spec_token_ids[last_req_index].clear() num_tokens = self.num_tokens[last_req_index] self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ From 87981eadda1f7f2efe0410290a5877c87a5141f4 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Tue, 18 Nov 2025 11:27:43 -0800 Subject: [PATCH 3/3] Preallocate spec tokens Signed-off-by: Jialin Ouyang --- vllm/v1/worker/gpu_input_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 113e0e0c5a38..4cd24523667f 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -251,7 +251,7 @@ def __init__( self.logitsprocs_need_output_token_ids = logitsprocs_need_output_token_ids # Store last speculative tokens for sampler. - self.spec_token_ids: list[list[int]] = [] + self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)] # This is updated each time the batch constituents change. self.sampling_metadata = self._make_sampling_metadata()