[V1][Spec Decode] Always use argmax for sampling draft tokens

[WoosukKwon]

This PR always samples draft tokens with argmax, regardless of the request’s temperature or other sampling parameters.
This will NOT affect the quality of the sampled outputs, but will affect (lower) the acceptance rate of the drafts, especially when the original temperature is high.

The reason behind this idea is that it’s tricky to handle the draft probability tensors efficiently.
If we use random sampling for draft tokens, we need to keep the draft probability tensors for rejection sampling.
However, because each of our scheduling step is (large model -> rejection sampling -> draft model), the draft probs tensors are not used immediately after they are created.
They are used when the corresponding draft tokens are scheduled, the timing of which is unpredictable.
Some of the draft tokens could be scheduled in the right next step, but others could be scheduled much later (e.g., because the request is preempted) or could be never scheduled (e.g., because the request has finished or been aborted).
This makes the management of the draft probs tensor as difficult as managing KV cache.

In contrast, if we use argmax sampling for draft tokens, we don't have to keep the draft prob tensor, which greatly simplifies the implementation.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[ShangmingCai]

It totally makes sense. The trade-off between implementation complexity and slightly lower acceptance rates (especially at high temperatures) seems reasonable given the current constraints. We can revisit this part later if community feedback indicates a strong need for higher acceptance rates under varying temperature settings.

[benchislett]

@WoosukKwon do you have any preliminary results for acceptance rate degradation? If the impact is minor (as I would expect), then I am comfortable with this change. If testing indicates a significant decrease in acceptance rates for high-temp sampling then we might want to reconsider.

[WoosukKwon]

@benchislett Good point. However, it's a bit tricky since we haven't implemented EAGLE with random sampling yet.

@luyuzhe111 Could you help with this, or perhaps share the script you used earlier to measure the acceptance length?

[luyuzhe111]

@WoosukKwon yea I can get the acceptance length with greedy drafting this week. the motivation behind this PR is also quite clear.

though if people want to contribute multi-draft spec dec in the future it will be impossible without draft probs.

[WoosukKwon]

@benchislett @luyuzhe111 I measured the acceptance lengths based on the EAGLE3 PR #16937:

Original temperature	0.0	0.6	0.7	0.8
Argmax	3.29	2.95	2.89	2.66
Same temp	3.29	3.30	3.25	3.11

[luyuzhe111]

@WoosukKwon thanks for sharing the results so fast! so is a 10% drop in AL acceptable? For chatbots I think 0.7 is a common temp choice?

[mergify]

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

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

[WoosukKwon]

so is a 10% drop in AL acceptable? For chatbots I think 0.7 is a common temp choice?

@luyuzhe111 I think this is definitely not acceptable in the long run, but I'm not sure whether we can use this PR as a temporary workaround.

[wwl2755]

I saw the PR #16077 has already aimed at part of this logic, maybe we could have the attempt from there?

[ekagra-ranjan]

FWIW, TRTLLM and TGI do not apply temp (greedy sampling) while sampling draft tokens even when the target model is using T!=0.

Is there any paper or resource which recommends applying the same temp to draft sampling as the target model for increased AL? I see this as an empirical evidence but wonder if there is some theorectical insights too.

[WoosukKwon]

Is there any paper or resource which recommends applying the same temp to draft sampling as the target model for increased AL?

@ekagra-ranjan Great question. I think it is backed by theory as well as experience. The draft token is accepted by rejection sampling when target_prob / draft_prob >= u where u is sampled from U(0, 1). Using argmax for draft tokens means making draft_prob always 1 (i.e., target_probs becomes the acceptance rate). Therefore, if the temperature for sampling target tokens is high, or if there are multiple good candidates for the next token, target_prob becomes small, so does the acceptance rate. If we use random sampling for draft tokens, draft_prob becomes smaller than 1, so the problem is mitigated generally.

[WoosukKwon]

@benchislett @luyuzhe111 @ekagra-ranjan @wwl2755 @ShangmingCai

Let me merge this PR first and get back to the acceptance rate issue later. The current main branch has a bug when using temp > 0, because it uses random sampling for draft tokens but does not consider draft probs in rejection sampling. This PR at least fixes this bug.

In the long run, I think the drop in the acceptance rate is unacceptable, so we should find out a better solution. This PR should be regarded as a band-aid solution.