[Bugfix][Scheduler] Don't pad spec decode up to `max_model_len` (#53962)
Nick Hill committed
Aug 27, 2026 at 08:30 UTC
76c0c6530e832c758f5502ca4ef37ae88a7ffa13
3 files changed
+57
-6
tests/v1/core/test_async_scheduler.py
+43
@@ -66,6 +66,49 @@ def test_stop_by_max_tokens(max_tokens: int):
66
assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens
67
68
69
+def test_no_spec_decode_padding_up_to_max_model_len():
70
+ """Uniform spec-decode padding must leave room for the sampled token.
71
+
72
+ Padding a single-token decode up to exactly max_model_len makes the
73
+ running-loop max_model_len cap negative on the next step (async
74
+ scheduling schedules it before the output that would stop the request),
75
+ and a negative num_scheduled_tokens crashes the model runner. The
76
+ request is scheduled un-padded instead.
77
+
78
+ Modelled on a disaggregated decode worker, where every request arrives
79
+ with all but the last prompt token already computed.
80
+ """
81
+ max_model_len = 32
82
+ num_spec = 1
83
+ scheduler = create_scheduler(
84
+ async_scheduling=True,
85
+ max_model_len=max_model_len,
86
+ num_speculative_tokens=num_spec,
87
+ speculative_method="ngram_gpu",
88
+ )
89
+
90
+ # A running decode, so padding's "something already scheduled and no
91
+ # prefill this step" precondition holds for the request below.
92
+ (filler,) = create_requests(num_requests=1, num_tokens=4, req_ids=["filler"])
93
+ scheduler.add_request(filler)
94
+ sched_output = scheduler.schedule()
95
+ scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
96
+
97
+ (request,) = create_requests(
98
+ num_requests=1, num_tokens=max_model_len - num_spec, req_ids=["boundary"]
99
+ )
100
+ request.num_computed_tokens = request.num_tokens - 1
101
+ scheduler.add_request(request)
102
+
103
+ sched_output = scheduler.schedule()
104
+ assert sched_output.num_scheduled_tokens["boundary"] == 1
105
+ assert "boundary" not in sched_output.scheduled_spec_decode_tokens
106
+ assert request.num_computed_tokens < max_model_len
107
+
108
+ scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output))
109
+ assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
110
+
111
+
112
def test_abort():
113
scheduler = create_scheduler(async_scheduling=True)
114
requests = create_requests(num_requests=10, max_tokens=20)
tests/v1/core/utils.py
+3
@@ -94,6 +94,9 @@ def create_scheduler(
94
dtype="float16",
95
seed=42,
96
skip_tokenizer_init=skip_tokenizer_init,
97
+ # The scheduler reads model_config.max_model_len, not the
98
+ # SchedulerConfig one, so both must agree.
99
+ max_model_len=max_model_len,
100
)
101
if use_ec_connector and ec_role == "ec_producer":
102
model_config.multimodal_config = MultiModalConfig()
vllm/v1/core/sched/scheduler.py
+11
-6
@@ -964,14 +964,19 @@ class Scheduler(SchedulerInterface):
964
and num_new_tokens == 1
965
and (scheduled_running_reqs and not prefill_scheduled)
966
):
967
- num_new_tokens = 1 + self.num_spec_tokens
967
+ padded_num_tokens = 1 + self.num_spec_tokens
968
+ # Pad only when there is room for the sampled token(s).
969
if (
969
- num_new_tokens > request_token_budget
970
- or num_computed_tokens + num_new_tokens > self.max_model_len
970
+ num_computed_tokens
971
+ + padded_num_tokens
972
+ + self.num_sampled_tokens_per_step
973
+ <= self.max_model_len
974
):
972
- # Prefer to not schedule than schedule un-padded here.
973
- break
974
- pad_spec_decode = True
975
+ if padded_num_tokens > request_token_budget:
976
+ # Prefer to not schedule than schedule un-padded.
977
+ break
978
+ num_new_tokens = padded_num_tokens
979
+ pad_spec_decode = True
980
981
threshold = self.scheduler_config.long_prefill_token_threshold
982
if 0 < threshold < num_new_tokens: