Skip to content

[V1][Spec Decode] Enable spec decode for top-p & top-k sampling #15063

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 13 commits into from
Mar 25, 2025
Merged

[V1][Spec Decode] Enable spec decode for top-p & top-k sampling #15063

merged 13 commits into from
Mar 25, 2025

Conversation

@WoosukKwon
Copy link
Collaborator

@WoosukKwon WoosukKwon commented Mar 18, 2025 •
edited by github-actions bot
Loading

This PR is to enable spec decoding for requests with top-p & top-k sampling.
It is implemented using apply_top_k_top_p to mask the logits of the target model.
While this is more expensive than FlashInfer's sorting-free sampling, I think it's good for the first step.

@WoosukKwon
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling
7c4f058 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
comment
4ce33c9 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
optional
e1d647e 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@github-actions
Copy link

github-actions bot commented Mar 18, 2025

👋 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.

🚀

@mergify mergify bot added the v1 label Mar 18, 2025
@WoosukKwon
minor
b1aeb04 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
minor
b1416a7 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
fix docstring
438e39e 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
Add example
0097a7b 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon WoosukKwon added the ready ONLY add when PR is ready to merge/full CI is needed label Mar 19, 2025
houseroad
houseroad reviewed Mar 19, 2025
View reviewed changes
houseroad
houseroad reviewed Mar 19, 2025
View reviewed changes
Copy link
Collaborator

@houseroad houseroad left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. Wondering how did we test it?

@WoosukKwon
Copy link
Collaborator Author

WoosukKwon commented Mar 19, 2025

Looks good. Wondering how did we test it?

Good point. I just wanted to get some initial feedback before adding tests. Will update the PR.

@WoosukKwon WoosukKwon added this to the v0.8.2 milestone Mar 24, 2025
@WoosukKwon
Add tests
5439c60 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
improve
2d10140 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon
Copy link
Collaborator Author

WoosukKwon commented Mar 24, 2025

@houseroad @LiuXiaoxuanPKU I've added the tests, and they're passing locally. Could you please review? I'd appreciate including this PR in v0.8.2 if possible.

LiuXiaoxuanPKU
LiuXiaoxuanPKU reviewed Mar 24, 2025
View reviewed changes
Copy link
Collaborator

@LiuXiaoxuanPKU LiuXiaoxuanPKU left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, just minor QQ about the test

tests/v1/sample/test_rejection_sampler.py Outdated
num_tokens = batch_size * num_draft_tokens

# Randomly create unmasked indices.
num_top_p_tokens = int(vocab_size * top_p)
Copy link
Collaborator

@LiuXiaoxuanPKU LiuXiaoxuanPKU Mar 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A bit confused by the definition of top_p sampling, should it be ' restricting the sampling to the set of most probable tokens with cumulative probability more than p'? Instead of sampling a fixed percentage of tokens like here.

Copy link
Collaborator Author

@WoosukKwon WoosukKwon Mar 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LiuXiaoxuanPKU Good catch. It only makes sense when int(vocab_size * top_p) tokens all have equal high logits (e.g., 100) while the others have -100. But definitely this is not general enough.

I've updated it to test top-p more precisely. Could you please take another look?

@WoosukKwon
fix top-p test
e66fab2 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
njhill
njhill approved these changes Mar 24, 2025
View reviewed changes
Copy link
Member

@njhill njhill left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@WoosukKwon
comment
5dfb42e 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
@WoosukKwon WoosukKwon merged commit ebcebee into main Mar 25, 2025
27 of 34 checks passed
@WoosukKwon WoosukKwon deleted the v1-spec-top-p branch March 25, 2025 00:16
@WoosukKwon WoosukKwon mentioned this pull request Mar 25, 2025
Merged
erictang000 pushed a commit to erictang000/vllm that referenced this pull request Mar 25, 2025
@WoosukKwon @erictang000
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
7425009 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
wrmedford pushed a commit to wrmedford/vllm that referenced this pull request Mar 26, 2025
@WoosukKwon @wrmedford
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
b238c2b 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Wes Medford <wryanmedford@gmail.com>
lulmer pushed a commit to lulmer/vllm that referenced this pull request Apr 7, 2025
@WoosukKwon @lulmer
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
c9aebbe 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Louis Ulmer <ulmerlouis@gmail.com>
@ckhordiasma ckhordiasma mentioned this pull request Apr 17, 2025
Closed
lk-chen pushed a commit to lk-chen/vllm that referenced this pull request Apr 29, 2025
@WoosukKwon @lk-chen
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
3b81918 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
shreyankg pushed a commit to shreyankg/vllm that referenced this pull request May 3, 2025
@WoosukKwon @shreyankg
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
be95135 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
RichardoMrMu pushed a commit to RichardoMrMu/vllm that referenced this pull request May 12, 2025
@WoosukKwon
[V1][Spec Decode] Enable spec decode for top-p & top-k sampling (vllm…
7eaf556 
…-project#15063)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@njhill njhill njhill approved these changes

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat

@ywang96 ywang96 Awaiting requested review from ywang96

@comaniac comaniac Awaiting requested review from comaniac

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat

@LiuXiaoxuanPKU LiuXiaoxuanPKU Awaiting requested review from LiuXiaoxuanPKU

@houseroad houseroad Awaiting requested review from houseroad houseroad is a code owner

Assignees

No one assigned

Labels

ready ONLY add when PR is ready to merge/full CI is needed v1

Projects

None yet

Milestone

v0.8.2

Development

Successfully merging this pull request may close these issues.

5 participants

@WoosukKwon @LiuXiaoxuanPKU @njhill @houseroad