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 10, 12:45 UTC
[BugFix] Fix spec decoding edge case bugs (#31944)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 7, 2026 at 23:31 UTC
287b37cda421f5e037e3b31fe330bfc4908faeaf
3 files changed
+49
-46
vllm/v1/core/sched/scheduler.py
+1
@@ -775,6 +775,7 @@ class Scheduler(SchedulerInterface):
775
self.encoder_cache_manager.free(request)
776
request.status = RequestStatus.PREEMPTED
777
request.num_computed_tokens = 0
778
+ request.spec_token_ids.clear()
779
request.num_preemptions += 1
780
if self.log_stats:
781
request.record_event(EngineCoreEventType.PREEMPTED, timestamp)
vllm/v1/worker/gpu_input_batch.py
+26
@@ -446,6 +446,32 @@ class InputBatch:
446
447
return req_index
448
449
+ def update_req_spec_token_ids(
450
+ self, request: CachedRequestState, scheduled_spec_tokens: dict[str, list[int]]
451
+ ) -> None:
452
+ req_id = request.req_id
453
+ req_index = self.req_id_to_index[req_id]
454
+ cur_spec_token_ids = self.spec_token_ids[req_index]
455
+ # When speculative decoding is used with structured output,
456
+ # the scheduler can drop draft tokens that do not
457
+ # conform to the schema. This can result in
458
+ # scheduler_output.scheduled_spec_decode_tokens being empty,
459
+ # even when speculative decoding is enabled.
460
+ cur_spec_token_ids.clear()
461
+ spec_token_ids = scheduled_spec_tokens.get(req_id, ())
462
+ num_spec_tokens = len(spec_token_ids)
463
+ request.prev_num_draft_len = num_spec_tokens
464
+ if not spec_token_ids:
465
+ return
466
+
467
+ # For async scheduling, token_ids_cpu assigned from
468
+ # spec_token_ids are placeholders and will be overwritten in
469
+ # _prepare_input_ids.
470
+ start_index = self.num_tokens_no_spec[req_index]
471
+ end_token_index = start_index + num_spec_tokens
472
+ self.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids
473
+ cur_spec_token_ids.extend(spec_token_ids)
474
+
475
def remove_request(self, req_id: str) -> int | None:
476
"""This method must always be followed by a call to condense().
477
vllm/v1/worker/gpu_model_runner.py
+22
-46
@@ -925,6 +925,7 @@ class GPUModelRunner(
925
# Update the states of the running/resumed requests.
926
is_last_rank = get_pp_group().is_last_rank
927
req_data = scheduler_output.scheduled_cached_reqs
928
+ scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
929
930
# Wait until valid_sampled_tokens_count is copied to cpu,
931
# then use it to update actual num_computed_tokens of each request.
@@ -938,20 +939,20 @@ class GPUModelRunner(
939
num_output_tokens = req_data.num_output_tokens[i]
940
req_index = self.input_batch.req_id_to_index.get(req_id)
941
941
- # prev_num_draft_len is used in async scheduling mode with
942
- # spec decode. it indicates if need to update num_computed_tokens
943
- # of the request. for example:
944
- # fist step: num_computed_tokens = 0, spec_tokens = [],
945
- # prev_num_draft_len = 0.
946
- # second step: num_computed_tokens = 100(prompt lenth),
947
- # spec_tokens = [a,b], prev_num_draft_len = 0.
948
- # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
949
- # prev_num_draft_len = 2.
950
- # num_computed_tokens in first step and second step does't contain
951
- # the spec tokens length, but in third step it contains the
952
- # spec tokens length. we only need to update num_computed_tokens
953
- # when prev_num_draft_len > 0.
954
- if req_state.prev_num_draft_len:
942
+ if req_state.prev_num_draft_len and self.use_async_scheduling:
943
+ # prev_num_draft_len is used in async scheduling mode with
944
+ # spec decode. it indicates if need to update num_computed_tokens
945
+ # of the request. for example:
946
+ # fist step: num_computed_tokens = 0, spec_tokens = [],
947
+ # prev_num_draft_len = 0.
948
+ # second step: num_computed_tokens = 100(prompt lenth),
949
+ # spec_tokens = [a,b], prev_num_draft_len = 0.
950
+ # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],
951
+ # prev_num_draft_len = 2.
952
+ # num_computed_tokens in first step and second step does't contain
953
+ # the spec tokens length, but in third step it contains the
954
+ # spec tokens length. we only need to update num_computed_tokens
955
+ # when prev_num_draft_len > 0.
956
if req_index is None:
957
req_state.prev_num_draft_len = 0
958
else:
@@ -1035,34 +1036,13 @@ class GPUModelRunner(
1036
self.input_batch.num_tokens_no_spec[req_index] = end_token_index
1037
1038
# Add spec_token_ids to token_ids_cpu.
1038
- spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
1039
- req_id, []
1040
- )
1041
- num_spec_tokens = len(spec_token_ids)
1042
- # For async scheduling, token_ids_cpu assigned from
1043
- # spec_token_ids are placeholders and will be overwritten in
1044
- # _prepare_input_ids.
1045
- if num_spec_tokens:
1046
- start_index = self.input_batch.num_tokens_no_spec[req_index]
1047
- end_token_index = start_index + num_spec_tokens
1048
- self.input_batch.token_ids_cpu[
1049
- req_index, start_index:end_token_index
1050
- ] = spec_token_ids
1039
+ self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens)
1040
1052
- # When speculative decoding is used with structured output,
1053
- # the scheduler can drop draft tokens that do not
1054
- # conform to the schema. This can result in
1055
- # scheduler_output.scheduled_spec_decode_tokens being empty,
1056
- # even when speculative decoding is enabled.
1057
- self.input_batch.spec_token_ids[req_index].clear()
1058
- self.input_batch.spec_token_ids[req_index].extend(spec_token_ids)
1059
-
1060
- if self.use_async_scheduling:
1061
- req_state.prev_num_draft_len = num_spec_tokens
1041
# Add the new or resumed requests to the persistent batch.
1042
# The smaller empty indices are filled first.
1043
for request in reqs_to_add:
1044
self.input_batch.add_request(request)
1045
+ self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens)
1046
1047
# Condense the batched states if there are gaps left by removed requests
1048
self.input_batch.condense()
@@ -1519,7 +1499,6 @@ class GPUModelRunner(
1499
# We will ignore the sampled tokens from the partial requests.
1500
# TODO: Support prompt logprobs.
1501
logits_indices = query_start_loc[1:] - 1
1522
- num_draft_tokens = None
1502
spec_decode_metadata = None
1503
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
1504
else:
@@ -1536,14 +1515,11 @@ class GPUModelRunner(
1515
) in scheduler_output.scheduled_spec_decode_tokens.items():
1516
req_idx = self.input_batch.req_id_to_index[req_id]
1517
num_draft_tokens[req_idx] = len(draft_token_ids)
1539
- num_decode_draft_tokens[req_idx] = (
1540
- len(draft_token_ids)
1541
- if (
1542
- self.input_batch.num_computed_tokens_cpu[req_idx]
1543
- >= self.input_batch.num_prompt_tokens[req_idx]
1544
- )
1545
- else -1
1546
- )
1518
+ if (
1519
+ self.input_batch.num_computed_tokens_cpu[req_idx]
1520
+ >= self.input_batch.num_prompt_tokens[req_idx]
1521
+ ):
1522
+ num_decode_draft_tokens[req_idx] = len(draft_token_ids)
1523
spec_decode_metadata = self._calc_spec_decode_metadata(
1524
num_draft_tokens, cu_num_tokens
1525
)