xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 20:45:49 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats 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) dim=-1)
# Get with the logprob of the prompt or sampled token. # 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) token_logprobs = logprobs.gather(-1, token_ids)
# Compute the ranks of the actual token. # Compute the ranks of the actual token.