frozenleaves/vllm
1
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2025-10-20 23:03:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix long dtype in topk sampling (#15049)

Browse Source
This commit is contained in:
Chujie Zheng
2025-03-19 06:57:31 +08:00
committed by GitHub
parent 72a8639b68
commit 027827cc1d
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
vllm/v1/sample/sampler.py
View File

@ -151,7 +151,7 @@ class Sampler(nn.Module):
dim=-1)
# Get with the logprob of the prompt or sampled token.
token_ids = token_ids.unsqueeze(-1)
token_ids = token_ids.unsqueeze(-1).to(torch.long)
token_logprobs = logprobs.gather(-1, token_ids)
# Compute the ranks of the actual token.