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, 16:00 UTC
[V1] Fix: make sure `k_index` is int64 for `apply_top_k_only` (#15907)
Signed-off-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Brayden Zhong committed
Apr 1, 2025 at 22:06 UTC
6efb195a6e02468001cc2a1ef818d55426da97c0
1 file changed
+1
-1
vllm/v1/sample/ops/topk_topp_sampler.py
+1
-1
@@ -200,7 +200,7 @@ def apply_top_k_only(
200
# topk.values tensor has shape [batch_size, max_top_k].
201
# Convert top k to 0-based index in range [0, max_top_k).
202
k_index = k.sub_(1).unsqueeze(1)
203
- top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index)
203
+ top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long())
204
# Handle non-topk rows.
205
top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf"))
206
logits.masked_fill_(logits < top_k_mask, -float("inf"))