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
[V1] Ensure using int64 for sampled token ids (#15065)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 18, 2025 at 23:52 UTC
05ccd0aa35581605017734c4ef36e058ddc58381
1 file changed
+8
-1
vllm/v1/sample/sampler.py
+8
-1
@@ -47,6 +47,11 @@ class Sampler(nn.Module):
47
logits = self.apply_penalties(logits, sampling_metadata)
48
# Sample the next token.
49
sampled = self.sample(logits, sampling_metadata)
50
+ # Convert sampled token ids to int64 (long) type to ensure compatibility
51
+ # with subsequent operations that may use these values as indices.
52
+ # This conversion is necessary because FlashInfer sampling operations
53
+ # return int32 (while PyTorch argmax and topk return int64).
54
+ sampled = sampled.long()
55
56
# Gather the logprobs of the topk and sampled token (if requested).
57
# Get logprobs and rank tensors (if requested)
@@ -139,19 +144,21 @@ class Sampler(nn.Module):
144
or sampled tokens (if sampled
145
logprobs); 1D token ID tensor
146
with (num tokens) elements
147
+ Must be int64.
148
149
Returns:
150
Top-k int indices tensor, (num tokens) x (num_logprobs + 1)
151
Top-k float logprobs tensor, (num tokens) x (num_logprobs + 1)
152
Sampled token rank tensor, (num tokens)
153
"""
154
+ assert token_ids.dtype == torch.int64
155
# Find the topK values.
156
topk_logprobs, topk_indices = torch.topk(logprobs,
157
num_logprobs,
158
dim=-1)
159
160
# Get with the logprob of the prompt or sampled token.
154
- token_ids = token_ids.unsqueeze(-1).to(torch.long)
161
+ token_ids = token_ids.unsqueeze(-1)
162
token_logprobs = logprobs.gather(-1, token_ids)
163
164
# Compute the ranks of the actual token.