From 373eb314af66559f9ef9d6b50bc2cb835012ef34 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Mon, 6 Jul 2026 20:50:38 +0700 Subject: [PATCH] [Bugfix][Core] Fix num_output_placeholders underflow with async scheduling + spec decode (#46066) Signed-off-by: Ting Sun Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/core/test_async_scheduler.py | 45 +++++++++++++++++++++++++++ tests/v1/core/utils.py | 8 ++++- vllm/v1/core/sched/scheduler.py | 8 +++-- 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 8253a8422e..9b6f645896 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -324,3 +324,48 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): assert request.status == RequestStatus.FINISHED_ERROR assert request.request_id not in scheduler.requests assert not scheduler.running + + +def test_no_placeholder_underflow_on_discarded_spec_frame(): + num_spec = 5 + scheduler = create_scheduler( + async_scheduling=True, + num_speculative_tokens=num_spec, + speculative_method="ngram_gpu", + ) + req = create_requests(num_requests=1, max_tokens=20)[0] + req.num_computed_tokens = req.num_tokens + scheduler.requests[req.request_id] = req + scheduler.running.append(req) + req.status = RequestStatus.RUNNING + + req.num_output_placeholders = 1 + req.async_tokens_to_discard = num_spec + computed_before = req.num_computed_tokens + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={req.request_id: num_spec + 1}, + total_num_scheduled_tokens=num_spec + 1, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={req.request_id: [10] * num_spec}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_runner_output = ModelRunnerOutput( + req_ids=[req.request_id], + req_id_to_index={req.request_id: 0}, + sampled_token_ids=[[999]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(scheduler_output, model_runner_output) + + assert req.num_output_placeholders == 1 + assert req.num_computed_tokens == computed_before + assert req.async_tokens_to_discard == num_spec - 1 + assert req.status == RequestStatus.RUNNING diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2450b23669..3a375ec757 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -54,6 +54,7 @@ def create_scheduler( block_size: int = 16, max_model_len: int | None = None, num_speculative_tokens: int | None = None, + speculative_method: str | None = None, skip_tokenizer_init: bool = False, async_scheduling: bool = False, pipeline_parallel_size: int = 1, @@ -126,9 +127,14 @@ def create_scheduler( speculative_config: SpeculativeConfig | None = None if num_speculative_tokens is not None: - speculative_config = SpeculativeConfig( + spec_kwargs: dict = dict( model="ngram", num_speculative_tokens=num_speculative_tokens ) + if speculative_method is not None: + spec_kwargs["method"] = speculative_method + spec_kwargs["prompt_lookup_max"] = num_speculative_tokens + spec_kwargs["prompt_lookup_min"] = 1 + speculative_config = SpeculativeConfig(**spec_kwargs) ec_transfer_config = ( ECTransferConfig( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9ce1d94ef3..f75bfae2b5 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1585,8 +1585,12 @@ class Scheduler(SchedulerInterface): scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) - if scheduled_spec_token_ids and ( - generated_token_ids or self.num_sampled_tokens_per_step == 0 + # Skip a stale frame still pending discard (async_tokens_to_discard + # > 0): its pre-reset rejection count would underflow the counters. + if ( + scheduled_spec_token_ids + and (generated_token_ids or self.num_sampled_tokens_per_step == 0) + and request.async_tokens_to_discard == 0 ): num_draft_tokens = len(scheduled_spec_token_ids) num_sampled = self.num_sampled_tokens_per_step