[Bugfix] Spec decode + structured output + spec model max len edge case

[andylolu2]

Purpose

Fixes #27210

Error:

  File "/home/ubuntu/vllm/vllm/v1/core/sched/scheduler.py", line 897, in get_grammar_bitmask
    bitmask = self.structured_output_manager.grammar_bitmask(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/vllm/vllm/v1/structured_output/__init__.py", line 263, in grammar_bitmask
    assert structured_output_request.grammar.accept_tokens(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

Cause:

The current implementation relies on scheduler.update_draft_token_ids to be called to reset previously-used spec_token_ids. However, scheduler.update_draft_token_ids is not guaranteed to be called after every step (currently it happens when the model_execute did not produce any new draft tokens, e.g. when input_fits_in_drafter is False). This results in stale spec_token_ids that don't necessarily follow the grammar after accepting the output tokens from the current step, leading to the above assertion error.

Proposed fix: Clear the spec token ids of a request immediately once it has been scheduled.

Test Plan

Extended current tests to cover this edge case.

Test Result

No more assertion errors.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request provides a solid fix for a critical bug involving speculative decoding and structured outputs, where stale speculative tokens could cause an assertion failure. The main change in vllm/v1/core/sched/scheduler.py to explicitly clear request.spec_token_ids after each scheduling step is correct and directly addresses the root cause. The subsequent simplification in update_draft_token_ids is a good cleanup that improves code clarity. Furthermore, the new test case in tests/v1/spec_decode/test_max_len.py is well-constructed to reproduce the specific edge case and prevent future regressions. The changes are of high quality and I see no issues.

[njhill]

Thanks a lot @andylolu2 for this fix! And for updating the test.

Does your test update now exercise the input_fits_in_drafter == False case? If so that's great, I had just noticed yesterday that there was no test added when that check was put in.

[andylolu2]

Thanks a lot @andylolu2 for this fix! And for updating the test.

Does your test update now exercise the input_fits_in_drafter == False case? If so that's great, I had just noticed yesterday that there was no test added when that check was put in.

@njhill I've updated the PR to address the comments, thanks for the review!

Yes the test now assertions the prompt_len is less than draft_max_len and prompt_len + output_len is greater than draft_max_len, so it triggers the input_fits_in_drafter == False case.

[njhill]

Thanks again @andylolu2