[Bugfix] Fall back to native sampling when FlashInfer cannot target the GPU (#48956)
Signed-off-by: martin el cheikh <martinelcheikh@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Martín el Cheikh committed
Sep 17, 2026 at 09:27 UTC
2e50824766ead3e30f53cba711e8c7ffb2bb6b2b
2 files changed
+76
-1
tests/v1/sample/test_topk_topp_sampler.py
+38
@@ -237,6 +237,44 @@ def test_flashinfer_sampler():
237
)
238
239
240
+@pytest.mark.skipif(
241
+ not FLASHINFER_TOPK_TOPP_SUPPORTED,
242
+ reason="Requires CUDA and flashinfer",
243
+)
244
+def test_flashinfer_sampler_fallback_when_jit_cannot_target_gpu(monkeypatch):
245
+ """FlashInfer swallows arch-detection errors (e.g. SM 12.x with a CUDA
246
+ toolkit older than 12.9), leaving an empty target-arch set that makes
247
+ every JIT spec fail with a misleading "requires sm75 or higher" error at
248
+ first use, killing the engine during startup profiling (issue #42393).
249
+ The sampler gate must detect this and fall back to native sampling, and
250
+ surface the real reason when the user explicitly opted in.
251
+ """
252
+ import flashinfer.jit.core
253
+ from flashinfer.compilation_context import CompilationContext
254
+
255
+ from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sampler_supported
256
+
257
+ def fake_check_cuda_arch():
258
+ raise RuntimeError("FlashInfer requires GPUs with sm75 or higher")
259
+
260
+ def fake_normalize_cuda_arch(major, minor):
261
+ raise RuntimeError("SM 12.x requires CUDA >= 12.9")
262
+
263
+ monkeypatch.setattr(flashinfer.jit.core, "check_cuda_arch", fake_check_cuda_arch)
264
+ monkeypatch.setattr(
265
+ CompilationContext,
266
+ "_normalize_cuda_arch",
267
+ staticmethod(fake_normalize_cuda_arch),
268
+ )
269
+
270
+ monkeypatch.delenv("VLLM_USE_FLASHINFER_SAMPLER", raising=False)
271
+ assert flashinfer_sampler_supported() is False
272
+
273
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "1")
274
+ with pytest.raises(RuntimeError, match="SM 12.x requires CUDA >= 12.9"):
275
+ flashinfer_sampler_supported()
276
+
277
+
278
# =============================================================================
279
# Triton kernel tests
280
# =============================================================================
vllm/v1/sample/ops/topk_topp_sampler.py
+38
-1
@@ -10,6 +10,7 @@ from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
11
from vllm.logger import init_logger
12
from vllm.platforms import CpuArchEnum, current_platform
13
+from vllm.platforms.interface import DeviceCapability
14
from vllm.triton_utils import HAS_TRITON
15
16
if HAS_TRITON:
@@ -41,12 +42,46 @@ def _skip_aiter_sampler_on_gfx1250() -> bool:
42
return on_gfx1250()
43
44
45
+def _flashinfer_jit_unsupported_reason(capability: DeviceCapability) -> str | None:
46
+ """Return why FlashInfer JIT codegen cannot target the current GPU, or
47
+ None if it can.
48
+
49
+ FlashInfer swallows arch-detection errors when building its compilation
50
+ context (e.g. SM 12.x with a CUDA toolkit older than 12.9), leaving an
51
+ empty target-arch set that makes every JIT spec fail with a misleading
52
+ "requires sm75 or higher" error at first use — killing the engine during
53
+ startup profiling (https://github.com/vllm-project/vllm/issues/42393).
54
+ """
55
+ try:
56
+ from flashinfer.jit.core import check_cuda_arch
57
+ except ImportError:
58
+ return None
59
+ try:
60
+ check_cuda_arch()
61
+ return None
62
+ except RuntimeError as e:
63
+ reason = str(e)
64
+ try:
65
+ # Re-derive the real error that FlashInfer swallowed during arch
66
+ # detection, e.g. "SM 12.x requires CUDA >= 12.9".
67
+ from flashinfer.compilation_context import CompilationContext
68
+
69
+ CompilationContext._normalize_cuda_arch(capability.major, capability.minor)
70
+ except RuntimeError as e:
71
+ reason = str(e)
72
+ except Exception:
73
+ # FlashInfer internals changed; keep the original error message.
74
+ pass
75
+ return reason
76
+
77
+
78
def flashinfer_sampler_supported() -> bool:
79
"""Decide whether FlashInfer's top-p/top-k sampler can be used.
80
81
Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER``
82
is 0, when the platform isn't CUDA, when the GPU's compute capability is
49
- unsupported. Raises ``RuntimeError`` if the user explicitly opted in
83
+ unsupported, or when FlashInfer cannot JIT-compile for the current
84
+ GPU/CUDA toolchain. Raises ``RuntimeError`` if the user explicitly opted in
85
via the env var but FlashInfer is unavailable.
86
87
Assumes flashinfer is installed, as guaranteed by ``requirements/cuda.txt``;
@@ -73,6 +108,8 @@ def flashinfer_sampler_supported() -> bool:
108
unsupported_reason = (
109
f"unsupported compute capability {capability.as_version_str()}"
110
)
111
+ else:
112
+ unsupported_reason = _flashinfer_jit_unsupported_reason(capability)
113
114
if unsupported_reason is None:
115
logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global")