[BugFix][Spec Decode] Use float64 for uniform_probs #23803
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Woosuk Kwon <[email protected]>
WoosukKwon
requested review from
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
mergify
bot
added
documentation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses a critical bug in the rejection sampler for speculative decoding. By changing the data type for uniform_probs from torch.float32 to torch.float64, it mitigates the risk of torch.rand producing an exact zero, which could lead to incorrect token acceptance. The change is well-justified and correctly implemented. The pull request also includes a fix in the spec_decode.py example to correctly handle prompt inputs. Both changes improve the correctness and robustness of the codebase. The changes look good to me.
youkaichao
approved these changes
Aug 28, 2025
zhewenl
pushed a commit
to zhewenl/vllm
that referenced
this pull request
Aug 28, 2025
Signed-off-by: Woosuk Kwon <[email protected]>
zhewenl
pushed a commit
to zhewenl/vllm
that referenced
this pull request
Sep 3, 2025
Signed-off-by: Woosuk Kwon <[email protected]>
eicherseiji
pushed a commit
to eicherseiji/vllm
that referenced
this pull request
Sep 9, 2025
Signed-off-by: Woosuk Kwon <[email protected]>
FeiDaLI
pushed a commit
to FeiDaLI/vllm
that referenced
this pull request
Sep 25, 2025
Signed-off-by: Woosuk Kwon <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Using float32 for
uniform_probscan be dangerous because exact 0 can be sampled with a non-negligible probability: pytorch/pytorch#16706This PR mitigates the issue by using float64 instead of float32