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 7, 22:45 UTC
Fix apply_top_k_top_p_triton called by non-cuda logits Tensor (#35030)
Signed-off-by: Xiao Li <ilx@meta.com>
Xiao Li committed
Feb 21, 2026 at 21:11 UTC
30132cd144af8876e7c0d2aac28cabaea3710254
1 file changed
+1
-1
vllm/v1/sample/ops/topk_topp_sampler.py
+1
-1
@@ -248,7 +248,7 @@ def apply_top_k_top_p(
248
if p is None and k is None:
249
return logits
250
251
- if HAS_TRITON and logits.shape[0] >= 8:
251
+ if HAS_TRITON and logits.shape[0] >= 8 and logits.is_cuda:
252
return apply_top_k_top_p_triton(logits, k, p)
253
254
# Use pytorch sort implementation for small batch sizes.