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 8, 22:45 UTC
[Core] Don't schedule spec tokens with prefill chunks (#33652)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 4, 2026 at 15:40 UTC
fa4e0fb028460cf5f4eb9cc90e206d0d6f35b026
5 files changed
+129
-29
tests/v1/core/test_scheduler.py
+94
@@ -945,6 +945,100 @@ def test_spec_decoding_stats_empty_output():
945
assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None
946
947
948
+def test_no_spec_tokens_scheduled_for_prefill_chunks():
949
+ """Test that draft tokens are ignored for prefill chunk requests.
950
+
951
+ When a request is being prefilled in chunks (chunked prefill), draft tokens
952
+ from `update_draft_token_ids` should be ignored until the prefill is complete.
953
+
954
+ The bug manifests when:
955
+ - A prefill chunk is scheduled
956
+ - Draft tokens are provided via update_draft_token_ids
957
+ - The next schedule has enough budget to include spec tokens
958
+
959
+ Without the fix, spec tokens would incorrectly be scheduled with the
960
+ remaining prefill tokens. With the fix, draft tokens are ignored for
961
+ prefill chunks.
962
+ """
963
+ num_spec_tokens = 3
964
+ # Use budget of 50, with 80 token prompt:
965
+ # - First chunk: 50 tokens
966
+ # - Second chunk: 30 remaining + potentially 3 spec tokens = 33
967
+ # Without fix: num_scheduled_spec_tokens = 33 + 50 - 80 = 3 (BUG!)
968
+ # With fix: spec_token_ids cleared, so no spec tokens scheduled
969
+ scheduler = create_scheduler(
970
+ num_speculative_tokens=num_spec_tokens,
971
+ max_num_batched_tokens=50,
972
+ enable_chunked_prefill=True,
973
+ )
974
+ requests = create_requests(num_requests=1, num_tokens=80)
975
+ req = requests[0]
976
+ scheduler.add_request(req)
977
+
978
+ # First schedule - prefill chunk (50 of 80 tokens)
979
+ output = scheduler.schedule()
980
+ assert len(output.scheduled_new_reqs) == 1
981
+ assert output.num_scheduled_tokens[req.request_id] == 50
982
+
983
+ # Update from output (no sampled token since still prefilling)
984
+ req_to_index = {req.request_id: 0}
985
+ model_runner_output = ModelRunnerOutput(
986
+ req_ids=[req.request_id],
987
+ req_id_to_index=req_to_index,
988
+ sampled_token_ids=[[]],
989
+ logprobs=None,
990
+ prompt_logprobs_dict={},
991
+ pooler_output=[],
992
+ )
993
+ scheduler.update_from_output(output, model_runner_output)
994
+
995
+ # Provide draft tokens while request is still in prefill.
996
+ # The fix ensures these are ignored for prefill chunks.
997
+ draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]])
998
+ scheduler.update_draft_token_ids(draft_token_ids)
999
+
1000
+ # Second schedule - remaining 30 tokens of prefill
1001
+ output = scheduler.schedule()
1002
+ # KEY ASSERTION: Should schedule exactly the remaining 30 prefill tokens,
1003
+ # NOT 33 (30 + 3 spec). Without the fix, this would be 33.
1004
+ assert output.num_scheduled_tokens[req.request_id] == 30, (
1005
+ f"Expected 30 tokens (remaining prefill only), "
1006
+ f"got {output.num_scheduled_tokens[req.request_id]}. "
1007
+ "Spec tokens should not be scheduled with prefill chunks."
1008
+ )
1009
+ # No spec tokens should be in the output
1010
+ assert req.request_id not in output.scheduled_spec_decode_tokens, (
1011
+ "Spec tokens should not be scheduled with prefill chunks"
1012
+ )
1013
+
1014
+ # Update from output with a sampled token (prefill complete)
1015
+ model_runner_output = ModelRunnerOutput(
1016
+ req_ids=[req.request_id],
1017
+ req_id_to_index=req_to_index,
1018
+ sampled_token_ids=[[42]],
1019
+ logprobs=None,
1020
+ prompt_logprobs_dict={},
1021
+ pooler_output=[],
1022
+ )
1023
+ scheduler.update_from_output(output, model_runner_output)
1024
+
1025
+ # Now provide draft tokens - should be accepted since prefill is complete
1026
+ draft_token_ids = DraftTokenIds([req.request_id], [[1, 2, 3]])
1027
+ scheduler.update_draft_token_ids(draft_token_ids)
1028
+
1029
+ # spec_token_ids SHOULD be set after prefill is complete
1030
+ assert req.spec_token_ids == [1, 2, 3], (
1031
+ f"spec_token_ids should be set after prefill, got {req.spec_token_ids}"
1032
+ )
1033
+
1034
+ # Third schedule - decode phase with spec tokens
1035
+ output = scheduler.schedule()
1036
+ # 1 new token + 3 spec tokens = 4
1037
+ assert output.num_scheduled_tokens[req.request_id] == 4
1038
+ assert req.request_id in output.scheduled_spec_decode_tokens
1039
+ assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens
1040
+
1041
+
1042
def _assert_right_scheduler_output(
1043
output: SchedulerOutput,
1044
num_requests: int,
vllm/v1/core/sched/async_scheduler.py
+10
-21
@@ -17,33 +17,22 @@ class AsyncScheduler(Scheduler):
17
18
def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
19
super()._update_after_schedule(scheduler_output)
20
- has_structured_output_requests = False
21
- pending_structured_output_tokens = False
20
spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens
21
for req_id in scheduler_output.num_scheduled_tokens:
22
request = self.requests[req_id]
25
- has_structured_output_requests |= request.use_structured_output
26
- pending_structured_output_tokens |= (
23
+ if request.is_prefill_chunk:
24
+ continue
25
+
26
+ scheduler_output.pending_structured_output_tokens |= (
27
request.use_structured_output and request.num_output_placeholders > 0
28
)
29
+ # The request will generate a new token plus num_spec_tokens
30
+ # in this scheduling step.
31
cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ()))
30
- if (
31
- request.num_computed_tokens
32
- == request.num_tokens
33
- + request.num_output_placeholders
34
- + cur_num_spec_tokens
35
- ):
36
- # The request will generate a new token plus num_spec_tokens
37
- # in this scheduling step.
38
- request.num_output_placeholders += 1 + cur_num_spec_tokens
39
- # Add placeholders for the new draft/spec tokens.
40
- # We will update the actual spec token ids in the worker process.
41
- request.spec_token_ids = self._spec_token_placeholders
42
-
43
- scheduler_output.has_structured_output_requests = has_structured_output_requests
44
- scheduler_output.pending_structured_output_tokens = (
45
- pending_structured_output_tokens
46
- )
32
+ request.num_output_placeholders += 1 + cur_num_spec_tokens
33
+ # Add placeholders for the new draft/spec tokens.
34
+ # We will update the actual spec token ids in the worker process.
35
+ request.spec_token_ids = self._spec_token_placeholders
36
37
def _update_request_with_output(
38
self, request: Request, new_token_ids: list[int]
vllm/v1/core/sched/scheduler.py
+12
@@ -912,6 +912,12 @@ class Scheduler(SchedulerInterface):
912
for req_id, num_scheduled_token in num_scheduled_tokens.items():
913
request = self.requests[req_id]
914
request.num_computed_tokens += num_scheduled_token
915
+ request.is_prefill_chunk = request.num_computed_tokens < (
916
+ request.num_tokens + request.num_output_placeholders
917
+ )
918
+ scheduler_output.has_structured_output_requests |= (
919
+ request.use_structured_output
920
+ )
921
922
# NOTE: _free_encoder_inputs relies on num_computed_tokens, which
923
# may be updated again in _update_from_output for speculative
@@ -1562,6 +1568,12 @@ class Scheduler(SchedulerInterface):
1568
# The request may have been finished. Skip.
1569
continue
1570
1571
+ if request.is_prefill_chunk:
1572
+ # Ignore draft tokens for prefill chunks.
1573
+ if request.spec_token_ids:
1574
+ request.spec_token_ids = []
1575
+ continue
1576
+
1577
# Add newly generated spec token ids to the request.
1578
if self.structured_output_manager.should_advance(request):
1579
metadata = request.structured_output_request
vllm/v1/request.py
+3
@@ -147,6 +147,9 @@ class Request:
147
# The number of tokens with prefix cache hits.
148
self.num_cached_tokens = -1
149
150
+ # True if this request is scheduled as a non-final prefill chunk.
151
+ self.is_prefill_chunk = False
152
+
153
# The number of NaNs in logits. A value greater than 0
154
# indicates that the output is corrupted
155
self.num_nans_in_logits = 0
vllm/v1/worker/gpu/spec_decode/utils.py
+10
-8
@@ -16,21 +16,21 @@ class DraftTokensHandler:
16
17
self.req_ids: list[str] = []
18
self.draft_tokens_np: np.ndarray | None = None
19
+ self.num_draft_tokens: int = 0
20
21
def set_draft_tokens(
22
self, input_batch: InputBatch, draft_tokens: torch.Tensor
23
) -> None:
24
+ self.req_ids = input_batch.req_ids
25
+ self.num_draft_tokens = draft_tokens.shape[1]
26
if not input_batch.has_structured_output_reqs:
27
# No draft token validation needs to be performed by
28
# the scheduler for this batch.
26
- if self.req_ids:
27
- self.req_ids = []
29
self.draft_tokens_np = None
30
return
31
32
# For spec decoding + structured outputs, we must transfer the
33
# draft tokens back to the scheduler for grammar validation.
33
- self.req_ids = input_batch.req_ids
34
current_stream = torch.cuda.current_stream(self.device)
35
self.copy_stream.wait_stream(current_stream)
36
with torch.cuda.stream(self.copy_stream):
@@ -38,8 +38,10 @@ class DraftTokensHandler:
38
self.copy_event.record()
39
40
def get_draft_tokens(self) -> DraftTokenIds | None:
41
- if self.draft_tokens_np is None:
42
- return None
43
-
44
- self.copy_event.synchronize()
45
- return DraftTokenIds(self.req_ids, self.draft_tokens_np.tolist())
41
+ if self.draft_tokens_np is not None:
42
+ self.copy_event.synchronize()
43
+ draft_token_ids = self.draft_tokens_np.tolist()
44
+ else:
45
+ # This case only happens when async scheduling is disabled.
46
+ draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
47
+ return DraftTokenIds(self.req_ids, draft_token_ids)