[Bugfix] Detect the CUDA toolkit the way FlashInfer does in has_flashinfer() (#58360)
Signed-off-by: Bowen Wang <bowen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Bowen Wang committed
Sep 29, 2026 at 12:14 UTC
1fde8ff7da1f96db844c4808559ef4a09eb4f680
2 files changed
+86
-6
tests/utils_/test_flashinfer_utils.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib.util
4
+from collections.abc import Iterator
5
+from pathlib import Path
6
+
7
+import pytest
8
+
9
+import vllm.utils.flashinfer as fi
10
+
11
+
12
+def _make_exe(path: Path) -> None:
13
+ path.parent.mkdir(parents=True, exist_ok=True)
14
+ path.touch(mode=0o755)
15
+
16
+
17
+@pytest.fixture
18
+def default_cuda_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
19
+ """Hide any real toolkit, keep ninja on PATH, redirect /usr/local/cuda."""
20
+ for var in ("CUDA_HOME", "CUDA_PATH", "FLASHINFER_NVCC"):
21
+ monkeypatch.delenv(var, raising=False)
22
+ _make_exe(tmp_path / "venv" / "bin" / "ninja")
23
+ monkeypatch.setenv("PATH", str(tmp_path / "venv" / "bin"))
24
+ home = tmp_path / "usr_local_cuda"
25
+ monkeypatch.setattr(fi, "_DEFAULT_CUDA_HOME", str(home))
26
+ return home
27
+
28
+
29
+@pytest.fixture(autouse=True)
30
+def flashinfer_without_cubin(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
31
+ """Pretend flashinfer-python is installed without flashinfer-cubin."""
32
+ real_find_spec = importlib.util.find_spec
33
+ monkeypatch.setattr(
34
+ importlib.util,
35
+ "find_spec",
36
+ lambda name, *args: (
37
+ object() if name == "flashinfer" else real_find_spec(name, *args)
38
+ ),
39
+ )
40
+ monkeypatch.setattr(fi, "has_flashinfer_cubin", lambda: False)
41
+ fi.has_flashinfer.cache_clear()
42
+ yield
43
+ fi.has_flashinfer.cache_clear()
44
+
45
+
46
+def test_has_flashinfer_finds_toolkit_off_path(default_cuda_home: Path):
47
+ """Bare-metal installs keep nvcc in /usr/local/cuda, off PATH, which is
48
+ where FlashInfer's JIT finds it."""
49
+ _make_exe(default_cuda_home / "bin" / "nvcc")
50
+ assert fi.has_flashinfer()
51
+
52
+
53
+def test_has_flashinfer_requires_ninja(
54
+ default_cuda_home: Path, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
55
+):
56
+ """FlashInfer runs `ninja` from PATH, which lacks the venv's bin directory
57
+ when vLLM is launched without activating the venv."""
58
+ _make_exe(tmp_path / "cuda" / "bin" / "nvcc")
59
+ monkeypatch.setenv("PATH", str(tmp_path / "cuda" / "bin"))
60
+ assert not fi.has_flashinfer()
vllm/utils/flashinfer.py
+26
-6
@@ -79,6 +79,23 @@ FLASHINFER_CUBINS_REPOSITORY = os.environ.get(
79
"https://edge.urm.nvidia.com/artifactory/sw-kernelinferencelibrary-public-generic-local/", # noqa: E501
80
)
81
82
+_DEFAULT_CUDA_HOME = "/usr/local/cuda"
83
+
84
+
85
+def _flashinfer_nvcc_path() -> str | None:
86
+ """Return the nvcc FlashInfer's JIT would run, or None if it is missing.
87
+
88
+ Mirrors ``flashinfer.jit.cpp_ext.get_cuda_path()`` without importing
89
+ FlashInfer, whose import initializes CUDA.
90
+ """
91
+ cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
92
+ if not cuda_home:
93
+ nvcc = shutil.which("nvcc")
94
+ cuda_home = (
95
+ os.path.dirname(os.path.dirname(nvcc)) if nvcc else _DEFAULT_CUDA_HOME
96
+ )
97
+ return shutil.which(os.path.join(cuda_home, "bin", "nvcc"))
98
+
99
100
@functools.cache
101
def has_flashinfer_cubin() -> bool:
@@ -99,12 +116,15 @@ def has_flashinfer() -> bool:
116
if importlib.util.find_spec("flashinfer") is None:
117
logger.debug_once("FlashInfer unavailable since package was not found")
118
return False
102
- # When not using flashinfer cubin,
103
- # Also check if nvcc is available since it's required to JIT compile flashinfer
104
- if not has_flashinfer_cubin() and shutil.which("nvcc") is None:
105
- logger.debug_once(
106
- "FlashInfer unavailable since nvcc was not found "
107
- "and not using pre-downloaded cubins"
119
+ # FlashInfer's JIT runs nvcc and `ninja` (from PATH).
120
+ if not has_flashinfer_cubin() and (
121
+ _flashinfer_nvcc_path() is None or shutil.which("ninja") is None
122
+ ):
123
+ logger.warning_once(
124
+ "FlashInfer kernels are disabled: flashinfer-cubin is not installed "
125
+ "and nvcc (CUDA_HOME, CUDA_PATH, PATH or /usr/local/cuda) or ninja "
126
+ "(PATH) is missing. Set CUDA_HOME to a CUDA toolkit and put ninja on "
127
+ "PATH, or run `flashinfer download-kernels`."
128
)
129
return False
130
return True