[Bugfix][Spec Decode] Don't pad a resumed decode request past max_model_len

[yifjiang]

Purpose

Fixes an engine-killing assertion when speculative decoding meets the context limit:

RuntimeError: Sampled token IDs exceed the max model length.
              Total number of tokens: 262145 > max_model_len: 262144

Seen twice in one day on a long-context production deployment (max_model_len 262144, MTP with
num_speculative_tokens: 1, prefix caching + CPU KV offload), under two different configurations.

Root cause

Two paths size a decode step, and they disagree by one token.

The running path reserves room for the token the step will sample —
scheduler.py#L593-L600:

num_new_tokens = min(
    num_new_tokens,
    self.max_model_len - request.num_computed_tokens - self.num_sampled_tokens_per_step,
)

The waiting path, padding a resumed decode request out to spec width for a full cudagraph,
does not — scheduler.py#L967-L974:

num_new_tokens = 1 + self.num_spec_tokens
if (num_new_tokens > request_token_budget
        or num_computed_tokens + num_new_tokens > self.max_model_len):   # permits ==
    break
pad_spec_decode = True

pad_spec_decode schedules real spec slots
(#L1161-L1164), so the step samples up to
1 + K tokens and the assertion in _bookkeeping_sync fires
(gpu_model_runner.py#L3903-L3909).

With C = num_computed_tokens and K = num_spec_tokens:

guard admits :  C + 1 + K <= max_model_len
runner needs :  C + 2 + K <= max_model_len

One apart, so it fires at exactly C == max_model_len - 1 - K and always overruns by one.

Reaching it needs speculative decoding, plus a resumed request holding a single uncomputed token
(preempted and brought back with its prefix recovered — routine with prefix caching or a KV-offload
tier), sitting at that one position.

Fix

Reserve num_sampled_tokens_per_step here too, so both paths agree.

On failure, fall back to an un-padded step rather than break. Unlike the token-budget case, this
condition never clears on a later step — num_computed_tokens cannot advance while the request is
unscheduled — so a bare guard fix would convert the crash into an indefinitely starved request. The
cost is one cudagraph miss on the final step of a request that is about to stop on length.

Test

test_resumed_decode_padded_to_spec_width_respects_max_model_len, parametrized over
num_spec_tokens ∈ {1, 3}. Scheduler-only, no model execution. Places the request at the derived
boundary and asserts both that the step cannot overrun and that it is still scheduled.

Measured before/after (max_model_len 2048):

            C     scheduled  spec   end_idx
before  K=1 2046      2        1     2049  > 2048   FAIL
before  K=3 2044      4        3     2049  > 2048   FAIL
after   K=1 2046      1        0     2048  = 2048   pass
after   K=3 2044      1        0     2046  < 2048   pass

Existing tests/v1/core/test_scheduler.py: 129 passed. Run on v0.26.0, where this hunk is identical
to main; I have not run main's full suite locally.

Separate from #52807 (offloading connector) — different subsystem, different failure.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aaron-seq]

In case you have not seen it, njhill opened #53962 for the same bug and says there that it replaces this PR and #50342. It has the ready label and CI running. Your diagnosis held up when I traced it, the two paths really do disagree by one token, so this was the right read. Might be worth checking that the un padded fallback in his version does what you wanted before this one gets closed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yifjiang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yifjiang]

Closing — superseded by #53962, which merged on 2026-08-27 and carries the same fix.

Confirmed on main: the padding guard now reserves the sampled token before padding, and falls back to an un-padded step rather than breaking out of the waiting loop.

Thanks @njhill for picking it up, and for folding in the create_scheduler(max_model_len=...) → ModelConfig fix — that trap cost me a bogus repro before I noticed the scheduler reads model_config.max_model_len, and fixing the helper is better than the runtime workaround I had.

One coverage note in case it is useful later: #53962's regression test asserts the async-scheduling path (negative num_scheduled_tokens → np.repeat). The same padding also trips _bookkeeping_sync's assert end_idx <= self.max_model_len on the sync path, which is the crash this PR was opened for, and that path is currently untested. I have a parametrized test over num_spec_tokens ∈ {1, 3} that places the request at max_model_len - 1 - K and asserts both no-overrun and not-starved; it failed on pre-fix main with 2049 > 2048. Happy to send it as a small test-only follow-up if you'd like it.