A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
fix long dtype in topk sampling (#15049)
Chujie Zheng committed
Mar 19, 2025 at 06:57 UTC
027827cc1d8d30d6e0108045ac9f9209a89dc59d
1 file changed
+1
-1
vllm/v1/sample/sampler.py
+1
-1
@@ -151,7 +151,7 @@ class Sampler(nn.Module):
151
dim=-1)
152
153
# Get with the logprob of the prompt or sampled token.
154
- token_ids = token_ids.unsqueeze(-1)
154
+ token_ids = token_ids.unsqueeze(-1).to(torch.long)
155
token_logprobs = logprobs.gather(-1, token_ids)
156
157
# Compute the ranks of the actual token.