[Bugfix][Core] Fix num_output_placeholders underflow with async scheduling + spec decode (#46066)
Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ting SUN committed
Jul 6, 2026 at 20:50 UTC
373eb314af66559f9ef9d6b50bc2cb835012ef34
3 files changed
+58
-3
tests/v1/core/test_async_scheduler.py
+45
@@ -324,3 +324,48 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
324
assert request.status == RequestStatus.FINISHED_ERROR
325
assert request.request_id not in scheduler.requests
326
assert not scheduler.running
327
+
328
+
329
+def test_no_placeholder_underflow_on_discarded_spec_frame():
330
+ num_spec = 5
331
+ scheduler = create_scheduler(
332
+ async_scheduling=True,
333
+ num_speculative_tokens=num_spec,
334
+ speculative_method="ngram_gpu",
335
+ )
336
+ req = create_requests(num_requests=1, max_tokens=20)[0]
337
+ req.num_computed_tokens = req.num_tokens
338
+ scheduler.requests[req.request_id] = req
339
+ scheduler.running.append(req)
340
+ req.status = RequestStatus.RUNNING
341
+
342
+ req.num_output_placeholders = 1
343
+ req.async_tokens_to_discard = num_spec
344
+ computed_before = req.num_computed_tokens
345
+
346
+ scheduler_output = SchedulerOutput(
347
+ scheduled_new_reqs=[],
348
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
349
+ num_scheduled_tokens={req.request_id: num_spec + 1},
350
+ total_num_scheduled_tokens=num_spec + 1,
351
+ scheduled_encoder_inputs={},
352
+ scheduled_spec_decode_tokens={req.request_id: [10] * num_spec},
353
+ num_common_prefix_blocks=[],
354
+ finished_req_ids=set(),
355
+ free_encoder_mm_hashes=[],
356
+ )
357
+ model_runner_output = ModelRunnerOutput(
358
+ req_ids=[req.request_id],
359
+ req_id_to_index={req.request_id: 0},
360
+ sampled_token_ids=[[999]],
361
+ logprobs=None,
362
+ prompt_logprobs_dict={},
363
+ pooler_output=[],
364
+ )
365
+
366
+ scheduler.update_from_output(scheduler_output, model_runner_output)
367
+
368
+ assert req.num_output_placeholders == 1
369
+ assert req.num_computed_tokens == computed_before
370
+ assert req.async_tokens_to_discard == num_spec - 1
371
+ assert req.status == RequestStatus.RUNNING
tests/v1/core/utils.py
+7
-1
@@ -54,6 +54,7 @@ def create_scheduler(
54
block_size: int = 16,
55
max_model_len: int | None = None,
56
num_speculative_tokens: int | None = None,
57
+ speculative_method: str | None = None,
58
skip_tokenizer_init: bool = False,
59
async_scheduling: bool = False,
60
pipeline_parallel_size: int = 1,
@@ -126,9 +127,14 @@ def create_scheduler(
127
128
speculative_config: SpeculativeConfig | None = None
129
if num_speculative_tokens is not None:
129
- speculative_config = SpeculativeConfig(
130
+ spec_kwargs: dict = dict(
131
model="ngram", num_speculative_tokens=num_speculative_tokens
132
)
133
+ if speculative_method is not None:
134
+ spec_kwargs["method"] = speculative_method
135
+ spec_kwargs["prompt_lookup_max"] = num_speculative_tokens
136
+ spec_kwargs["prompt_lookup_min"] = 1
137
+ speculative_config = SpeculativeConfig(**spec_kwargs)
138
139
ec_transfer_config = (
140
ECTransferConfig(
vllm/v1/core/sched/scheduler.py
+6
-2
@@ -1585,8 +1585,12 @@ class Scheduler(SchedulerInterface):
1585
scheduled_spec_token_ids = (
1586
scheduler_output.scheduled_spec_decode_tokens.get(req_id)
1587
)
1588
- if scheduled_spec_token_ids and (
1589
- generated_token_ids or self.num_sampled_tokens_per_step == 0
1588
+ # Skip a stale frame still pending discard (async_tokens_to_discard
1589
+ # > 0): its pre-reset rejection count would underflow the counters.
1590
+ if (
1591
+ scheduled_spec_token_ids
1592
+ and (generated_token_ids or self.num_sampled_tokens_per_step == 0)
1593
+ and request.async_tokens_to_discard == 0
1594
):
1595
num_draft_tokens = len(scheduled_spec_token_ids)
1596
num_sampled = self.num_sampled_tokens_per_step