biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[BugFix][Spec Decode] Fix spec token ids in model runner (#20530)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon
2025-07-06 12:44:52 -07:00
committed by GitHub
parent 9fb52e523a
commit 9528e3a05e
1 changed files with 12 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
vllm/v1/worker/gpu_model_runner.py
View File

@@ -528,18 +528,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
start_token_index:end_token_index] = new_token_ids start_token_index:end_token_index] = new_token_ids
self.input_batch.num_tokens_no_spec[ self.input_batch.num_tokens_no_spec[
req_index] = end_token_index req_index] = end_token_index
self.input_batch.num_tokens[req_index] = end_token_index
# Add spec_token_ids to token_ids_cpu. # Add spec_token_ids to token_ids_cpu.
spec_token_ids = ( spec_token_ids = (
scheduler_output.scheduled_spec_decode_tokens.get( scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
req_id, ()))
if spec_token_ids: if spec_token_ids:
start_index = end_token_index num_spec_tokens = len(spec_token_ids)
end_token_index += len(spec_token_ids) start_index = self.input_batch.num_tokens_no_spec[req_index]
end_token_index = start_index + num_spec_tokens
self.input_batch.token_ids_cpu[ self.input_batch.token_ids_cpu[
req_index, req_index, start_index:end_token_index] = spec_token_ids
start_index:end_token_index] = spec_token_ids
# NOTE(woosuk): `num_tokens` here may include spec tokens. # NOTE(woosuk): `num_tokens` here may include spec tokens.
self.input_batch.num_tokens[req_index] = end_token_index self.input_batch.num_tokens[req_index] += num_spec_tokens
# Add the new or resumed requests to the persistent batch. # Add the new or resumed requests to the persistent batch.
# The smaller empty indices are filled first. # The smaller empty indices are filled first.