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, 19:00 UTC
remove cuda check in `top_k_top_p_triton` kernel (#35011)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Feb 25, 2026 at 14:22 UTC
92510edc325c855848653c8864f0805eb7fbb022
2 files changed
+3
-4
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 and logits.is_cuda:
251
+ if HAS_TRITON and logits.shape[0] >= 8:
252
return apply_top_k_top_p_triton(logits, k, p)
253
254
# Use pytorch sort implementation for small batch sizes.
vllm/v1/sample/ops/topk_topp_triton.py
+2
-3
@@ -967,7 +967,6 @@ def apply_top_k_top_p_triton(
967
"""
968
assert logits.ndim == 2
969
assert logits.dtype == torch.float32
970
- assert logits.is_cuda
970
971
batch_size, vocab_size = logits.shape
972
@@ -978,13 +977,13 @@ def apply_top_k_top_p_triton(
977
return logits
978
979
if k is not None:
981
- assert k.ndim == 1 and k.shape[0] == batch_size and k.is_cuda
980
+ assert k.ndim == 1 and k.shape[0] == batch_size
981
k_ptr = k.to(torch.int32)
982
else:
983
k_ptr = logits # Dummy pointer (won't be read)
984
985
if p is not None:
987
- assert p.ndim == 1 and p.shape[0] == batch_size and p.is_cuda
986
+ assert p.ndim == 1 and p.shape[0] == batch_size
987
p_ptr = p.to(torch.float32)
988
else:
989
p_ptr = logits # Dummy pointer (won't be read)