[Bug][Spec Decode] Fix AR metrics when drafting is skipped

[benchislett]

Purpose

FIX #34734.

Not the cleanest fix, open to other approaches. But this one does seem to work

Testing

Ran Llama 3.1 8B-Instruct with EAGLE3 and got AL 1.00 for prompts of size 2k. Now, reports no drafted tokens as expected.

Will fix CI failures as needed, but shouldn't affect anything other than statistics collection.

[gemini-code-assist]

Code Review

The pull request addresses a bug where acceptance rate (AR) metrics for speculative decoding were incorrectly calculated when drafting was skipped. The fix introduces a new boolean flag _prev_step_drafting_was_skipped in the Scheduler class and a drafting_was_skipped field in ModelRunnerOutput to track this state. This flag is then used to conditionally exclude dummy/placeholder draft tokens from AR statistics. The changes appear to correctly handle the scenario described in the bug report, ensuring more accurate metrics for speculative decoding.

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!

[mergify]

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

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

[CarrotSwordsman]

@benchislett I duplicated this work in #54054 before finding this PR — my fault for not checking open PRs against #34734 properly. I have closed mine. One finding from it that may be worth folding in here, since this PR is currently stale with conflicts:

The async-scheduling path needs separate handling. Consuming the flag in update_from_output covers synchronous scheduling, but with async_scheduling=True the draft tokens for the step are already part of the in-flight batch by the time the flag arrives, so they cannot be un-scheduled. In update_draft_token_ids_in_output they instead have to go down the existing invalid-draft channel — padded to -1 and recorded in num_invalid_spec_tokens — which is the same path grammar-invalidated drafts already take, and which make_spec_decoding_stats subtracts from num_draft_tokens.

Also worth noting: there is a residual case in the stats function itself. make_spec_decoding_stats tests if not num_draft_tokens: return before subtracting num_invalid_spec_tokens, so on the async path num_drafts still increments even when every drafted token is invalid. num_draft_tokens ends up 0 so the misleading 0.00% AR is gone, but the drafts counter is still inflated. That ordering predates both our PRs and affects structured-output invalidation too, so it may deserve its own fix rather than being bundled here.

Happy to help rebase this one or test a revision if useful — though I only have CPU-level scheduler testing available, no spec-decode checkpoint.

[seongyun1104]

This is still needed on main, and the follow-up you deferred has since become a two-way race that neither side has connected back to here. Four things, in case any of them change what you want to do with this branch.

1. Not obsolete, but your touchpoint moved. The predicate is still there and still batch-wide, so #34734 reproduces as written — the stale-bot close in June was not the code moving out from under it. It has been extracted into a helper since you wrote this, which is where your rebase will land:

# vllm/v1/worker/gpu_model_runner.py, main fa1b3b1922
:4550   def _input_fits_in_drafter(self, common_attn_metadata) -> bool:
:4553       if common_attn_metadata is None:
:4554           return False
:4557       num_drafter_query_tokens = self.num_spec_tokens + (
:4558           1 if self.speculative_config.use_dflash() else 0
:4559       )
:4560       return (
:4561           common_attn_metadata.max_seq_len + num_drafter_query_tokens
:4562           <= self.effective_drafter_max_model_len
:4563       )

called at :4644. Your diff sets drafting_was_skipped inline beside the old expression, so that hunk will not apply as-is; the natural spot now is at the call site rather than inside the predicate.

2. The residual CarrotSwordsman flagged is now being fixed by two other PRs, and neither cites #34734 or this branch. They are #56137 and #56195, both aimed at exactly the ordering you were told might deserve its own fix — make_spec_decoding_stats testing num_draft_tokens before subtracting num_invalid_spec_tokens. I checked: neither PR body references either number, so as far as I can tell they do not know this thread exists.

3. One of those two approaches is unsafe, and it is worth knowing which. #56195 moves the grammar adjustment ahead of the guard, so the function can return None for a request whose adjusted count is zero. On main that is not a local skip:

# vllm/v1/core/sched/scheduler.py, main fa1b3b1922
:1901   spec_decoding_stats: SpecDecodingStats | None = None      # once per update_from_output, outside the loop
:1997   spec_decoding_stats = self.make_spec_decoding_stats(      # inside the loop, return value assigned back
:1998       spec_decoding_stats, ...
:2294   ...                                                        # consumed once, outside the loop

Because the return value is assigned back into the accumulator, an early return None in the middle of the loop discards every request already folded in for that step, not just the one being skipped. It is order-dependent and silent. #56137 keeps the guard on the raw count and subtracts afterwards, so it does not have this path. I said as much on #56195; repeating it here only because this branch is where the question originated.

Your own gate sits outside make_spec_decoding_stats rather than inside it, so it does not have that problem — but if either of those lands first, the two gates will be sitting on the same call site from opposite sides.

4. A composition note, and my stake in it. Your fix suppresses acceptance stats for a step where the drafter was skipped, which is right for the AR number — input_fits_in_drafter is computed from the batch-wide max_seq_len, so the skip really is whole-step and the whole-step suppression matches it. The side effect is that the set of steps producing no SpecDecodingStats at all gets larger: after this, both "the drafter was skipped for length" and "the schedule selected K=0" are silent in the same way, and neither is distinguishable from speculative decoding simply not running.

I have an open PR on that second case (#54748, a step counter keyed on the selected K), so I am not a neutral party here. I am not asking you to adopt or reference it — only flagging that the two changes push the same surface in opposite directions, and you may want to know that before rebasing.

No action needed from me; happy to be told any of this is already accounted for.