diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 967f0857a1..7f256f025e 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -66,6 +66,49 @@ def test_stop_by_max_tokens(max_tokens: int): assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens +def test_no_spec_decode_padding_up_to_max_model_len(): + """Uniform spec-decode padding must leave room for the sampled token. + + Padding a single-token decode up to exactly max_model_len makes the + running-loop max_model_len cap negative on the next step (async + scheduling schedules it before the output that would stop the request), + and a negative num_scheduled_tokens crashes the model runner. The + request is scheduled un-padded instead. + + Modelled on a disaggregated decode worker, where every request arrives + with all but the last prompt token already computed. + """ + max_model_len = 32 + num_spec = 1 + scheduler = create_scheduler( + async_scheduling=True, + max_model_len=max_model_len, + num_speculative_tokens=num_spec, + speculative_method="ngram_gpu", + ) + + # A running decode, so padding's "something already scheduled and no + # prefill this step" precondition holds for the request below. + (filler,) = create_requests(num_requests=1, num_tokens=4, req_ids=["filler"]) + scheduler.add_request(filler) + sched_output = scheduler.schedule() + scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) + + (request,) = create_requests( + num_requests=1, num_tokens=max_model_len - num_spec, req_ids=["boundary"] + ) + request.num_computed_tokens = request.num_tokens - 1 + scheduler.add_request(request) + + sched_output = scheduler.schedule() + assert sched_output.num_scheduled_tokens["boundary"] == 1 + assert "boundary" not in sched_output.scheduled_spec_decode_tokens + assert request.num_computed_tokens < max_model_len + + scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED + + def test_abort(): scheduler = create_scheduler(async_scheduling=True) requests = create_requests(num_requests=10, max_tokens=20) diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 802326d306..d3090e2dda 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -94,6 +94,9 @@ def create_scheduler( dtype="float16", seed=42, skip_tokenizer_init=skip_tokenizer_init, + # The scheduler reads model_config.max_model_len, not the + # SchedulerConfig one, so both must agree. + max_model_len=max_model_len, ) if use_ec_connector and ec_role == "ec_producer": model_config.multimodal_config = MultiModalConfig() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c373d5cf5c..51f75a63d6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -964,14 +964,19 @@ class Scheduler(SchedulerInterface): and num_new_tokens == 1 and (scheduled_running_reqs and not prefill_scheduled) ): - num_new_tokens = 1 + self.num_spec_tokens + padded_num_tokens = 1 + self.num_spec_tokens + # Pad only when there is room for the sampled token(s). if ( - num_new_tokens > request_token_budget - or num_computed_tokens + num_new_tokens > self.max_model_len + num_computed_tokens + + padded_num_tokens + + self.num_sampled_tokens_per_step + <= self.max_model_len ): - # Prefer to not schedule than schedule un-padded here. - break - pad_spec_decode = True + if padded_num_tokens > request_token_budget: + # Prefer to not schedule than schedule un-padded. + break + num_new_tokens = padded_num_tokens + pad_spec_decode = True threshold = self.scheduler_config.long_prefill_token_threshold if 0 < threshold < num_new_tokens: