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 15, 03:30 UTC
[BugFix][Spec Decode] Fix spec token ids in model runner (#20530)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jul 6, 2025 at 12:44 UTC
9528e3a05e4264e7477bfe4e30a4e1d8fa54e1dc
1 file changed
+12
-11
vllm/v1/worker/gpu_model_runner.py
+12
-11
@@ -528,19 +528,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
528
start_token_index:end_token_index] = new_token_ids
529
self.input_batch.num_tokens_no_spec[
530
req_index] = end_token_index
531
- # Add spec_token_ids to token_ids_cpu.
532
- spec_token_ids = (
533
- scheduler_output.scheduled_spec_decode_tokens.get(
534
- req_id, ()))
535
- if spec_token_ids:
536
- start_index = end_token_index
537
- end_token_index += len(spec_token_ids)
538
- self.input_batch.token_ids_cpu[
539
- req_index,
540
- start_index:end_token_index] = spec_token_ids
541
- # NOTE(woosuk): `num_tokens` here may include spec tokens.
531
self.input_batch.num_tokens[req_index] = end_token_index
532
533
+ # Add spec_token_ids to token_ids_cpu.
534
+ spec_token_ids = (
535
+ scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
536
+ if spec_token_ids:
537
+ num_spec_tokens = len(spec_token_ids)
538
+ start_index = self.input_batch.num_tokens_no_spec[req_index]
539
+ end_token_index = start_index + num_spec_tokens
540
+ self.input_batch.token_ids_cpu[
541
+ req_index, start_index:end_token_index] = spec_token_ids
542
+ # NOTE(woosuk): `num_tokens` here may include spec tokens.
543
+ self.input_batch.num_tokens[req_index] += num_spec_tokens
544
+
545
# Add the new or resumed requests to the persistent batch.
546
# The smaller empty indices are filled first.
547
for req_id in req_ids_to_add: