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 12, 09:45 UTC
Guard FlashInfer sampler using the same check as FlashInfer attention backend (#29415)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Harry Mellor committed
Nov 28, 2025 at 16:34 UTC
9eec282cb5a6fe0b39449476dc4d14da0516984c
1 file changed
+10
vllm/v1/sample/ops/topk_topp_sampler.py
+10
@@ -33,6 +33,16 @@ class TopKTopPSampler(nn.Module):
33
and current_platform.is_cuda()
34
):
35
if envs.VLLM_USE_FLASHINFER_SAMPLER:
36
+ from vllm.v1.attention.backends.flashinfer import FlashInferBackend
37
+
38
+ capability = current_platform.get_device_capability()
39
+ assert capability is not None
40
+ if not FlashInferBackend.supports_compute_capability(capability):
41
+ capability_str = capability.as_version_str()
42
+ raise RuntimeError(
43
+ "FlashInfer does not support compute capability "
44
+ f"{capability_str}, unset VLLM_USE_FLASHINFER_SAMPLER=1."
45
+ )
46
# Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1.
47
logger.info_once(
48
"Using FlashInfer for top-p & top-k sampling.",