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 6, 14:30 UTC
[Bugfix](xpu): prevent “selected index k out of range” in TP decode path (#37259)
Signed-off-by: zhenzhao <zhenzhao@habana.ai>
zhao, zhenhui committed
Mar 17, 2026 at 19:14 UTC
4af9ed21cba9e4bb85cd7cc124aa6f23cd0ae9a5
1 file changed
+2
-1
vllm/_xpu_ops.py
+2
-1
@@ -426,7 +426,8 @@ class xpu_ops:
426
mask = positions <= index_end_pos
427
# mask: [B * N, L]
428
logits = logits.masked_fill(~mask, float("-inf"))
429
- topk_indices = logits.topk(topk_tokens, dim=-1)[1].to(torch.int32) # [B * N, K]
429
+ real_topk = min(topk_tokens, logits.shape[-1])
430
+ topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32) # [B * N, K]
431
# ensure we don't set indices for the top k
432
# that is out of range(masked already)
433
# this will happen if context length is shorter than K