From 2e50824766ead3e30f53cba711e8c7ffb2bb6b2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mart=C3=ADn=20el=20Cheikh?= <157884230+melcheikh@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:27:54 -0300 Subject: [PATCH] [Bugfix] Fall back to native sampling when FlashInfer cannot target the GPU (#48956) Signed-off-by: martin el cheikh Co-authored-by: Claude --- tests/v1/sample/test_topk_topp_sampler.py | 38 ++++++++++++++++++++++ vllm/v1/sample/ops/topk_topp_sampler.py | 39 ++++++++++++++++++++++- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 0b570f49d9..8ff16e5599 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -237,6 +237,44 @@ def test_flashinfer_sampler(): ) +@pytest.mark.skipif( + not FLASHINFER_TOPK_TOPP_SUPPORTED, + reason="Requires CUDA and flashinfer", +) +def test_flashinfer_sampler_fallback_when_jit_cannot_target_gpu(monkeypatch): + """FlashInfer swallows arch-detection errors (e.g. SM 12.x with a CUDA + toolkit older than 12.9), leaving an empty target-arch set that makes + every JIT spec fail with a misleading "requires sm75 or higher" error at + first use, killing the engine during startup profiling (issue #42393). + The sampler gate must detect this and fall back to native sampling, and + surface the real reason when the user explicitly opted in. + """ + import flashinfer.jit.core + from flashinfer.compilation_context import CompilationContext + + from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sampler_supported + + def fake_check_cuda_arch(): + raise RuntimeError("FlashInfer requires GPUs with sm75 or higher") + + def fake_normalize_cuda_arch(major, minor): + raise RuntimeError("SM 12.x requires CUDA >= 12.9") + + monkeypatch.setattr(flashinfer.jit.core, "check_cuda_arch", fake_check_cuda_arch) + monkeypatch.setattr( + CompilationContext, + "_normalize_cuda_arch", + staticmethod(fake_normalize_cuda_arch), + ) + + monkeypatch.delenv("VLLM_USE_FLASHINFER_SAMPLER", raising=False) + assert flashinfer_sampler_supported() is False + + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "1") + with pytest.raises(RuntimeError, match="SM 12.x requires CUDA >= 12.9"): + flashinfer_sampler_supported() + + # ============================================================================= # Triton kernel tests # ============================================================================= diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 06b3fc3423..477e99c2fa 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -10,6 +10,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform +from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import HAS_TRITON if HAS_TRITON: @@ -41,12 +42,46 @@ def _skip_aiter_sampler_on_gfx1250() -> bool: return on_gfx1250() +def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None: + """Return why FlashInfer JIT codegen cannot target the current GPU, or + None if it can. + + FlashInfer swallows arch-detection errors when building its compilation + context (e.g. SM 12.x with a CUDA toolkit older than 12.9), leaving an + empty target-arch set that makes every JIT spec fail with a misleading + "requires sm75 or higher" error at first use — killing the engine during + startup profiling (https://github.com/vllm-project/vllm/issues/42393). + """ + try: + from flashinfer.jit.core import check_cuda_arch + except ImportError: + return None + try: + check_cuda_arch() + return None + except RuntimeError as e: + reason = str(e) + try: + # Re-derive the real error that FlashInfer swallowed during arch + # detection, e.g. "SM 12.x requires CUDA >= 12.9". + from flashinfer.compilation_context import CompilationContext + + CompilationContext._normalize_cuda_arch(capability.major, capability.minor) + except RuntimeError as e: + reason = str(e) + except Exception: + # FlashInfer internals changed; keep the original error message. + pass + return reason + + def flashinfer_sampler_supported() -> bool: """Decide whether FlashInfer's top-p/top-k sampler can be used. Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER`` is 0, when the platform isn't CUDA, when the GPU's compute capability is - unsupported. Raises ``RuntimeError`` if the user explicitly opted in + unsupported, or when FlashInfer cannot JIT-compile for the current + GPU/CUDA toolchain. Raises ``RuntimeError`` if the user explicitly opted in via the env var but FlashInfer is unavailable. Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``; @@ -73,6 +108,8 @@ def flashinfer_sampler_supported() -> bool: unsupported_reason = ( f"unsupported compute capability {capability.as_version_str()}" ) + else: + unsupported_reason = _flashinfer_jit_unsupported_reason(capability) if unsupported_reason is None: logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global")