diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 5b2d130b0ea42..c9229e788b6bf 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -33,6 +33,16 @@ class TopKTopPSampler(nn.Module): and current_platform.is_cuda() ): if envs.VLLM_USE_FLASHINFER_SAMPLER: + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + capability = current_platform.get_device_capability() + assert capability is not None + if not FlashInferBackend.supports_compute_capability(capability): + capability_str = capability.as_version_str() + raise RuntimeError( + "FlashInfer does not support compute capability " + f"{capability_str}, unset VLLM_USE_FLASHINFER_SAMPLER=1." + ) # Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1. logger.info_once( "Using FlashInfer for top-p & top-k sampling.",