[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>

tests/utils_/test_flashinfer_utils.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib.util
+from collections.abc import Iterator
+from pathlib import Path
+
+import pytest
+
+import vllm.utils.flashinfer as fi
+
+
+def _make_exe(path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.touch(mode=0o755)
+
+
+@pytest.fixture
+def default_cuda_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
+    """Hide any real toolkit, keep ninja on PATH, redirect /usr/local/cuda."""
+    for var in ("CUDA_HOME", "CUDA_PATH", "FLASHINFER_NVCC"):
+        monkeypatch.delenv(var, raising=False)
+    _make_exe(tmp_path / "venv" / "bin" / "ninja")
+    monkeypatch.setenv("PATH", str(tmp_path / "venv" / "bin"))
+    home = tmp_path / "usr_local_cuda"
+    monkeypatch.setattr(fi, "_DEFAULT_CUDA_HOME", str(home))
+    return home
+
+
+@pytest.fixture(autouse=True)
+def flashinfer_without_cubin(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    """Pretend flashinfer-python is installed without flashinfer-cubin."""
+    real_find_spec = importlib.util.find_spec
+    monkeypatch.setattr(
+        importlib.util,
+        "find_spec",
+        lambda name, *args: (
+            object() if name == "flashinfer" else real_find_spec(name, *args)
+        ),
+    )
+    monkeypatch.setattr(fi, "has_flashinfer_cubin", lambda: False)
+    fi.has_flashinfer.cache_clear()
+    yield
+    fi.has_flashinfer.cache_clear()
+
+
+def test_has_flashinfer_finds_toolkit_off_path(default_cuda_home: Path):
+    """Bare-metal installs keep nvcc in /usr/local/cuda, off PATH, which is
+    where FlashInfer's JIT finds it."""
+    _make_exe(default_cuda_home / "bin" / "nvcc")
+    assert fi.has_flashinfer()
+
+
+def test_has_flashinfer_requires_ninja(
+    default_cuda_home: Path, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+):
+    """FlashInfer runs `ninja` from PATH, which lacks the venv's bin directory
+    when vLLM is launched without activating the venv."""
+    _make_exe(tmp_path / "cuda" / "bin" / "nvcc")
+    monkeypatch.setenv("PATH", str(tmp_path / "cuda" / "bin"))
+    assert not fi.has_flashinfer()

vllm/utils/flashinfer.py

@@ -79,6 +79,23 @@ FLASHINFER_CUBINS_REPOSITORY = os.environ.get(
     "https://edge.urm.nvidia.com/artifactory/sw-kernelinferencelibrary-public-generic-local/",  # noqa: E501
 )
 
+_DEFAULT_CUDA_HOME = "/usr/local/cuda"
+
+
+def _flashinfer_nvcc_path() -> str | None:
+    """Return the nvcc FlashInfer's JIT would run, or None if it is missing.
+
+    Mirrors ``flashinfer.jit.cpp_ext.get_cuda_path()`` without importing
+    FlashInfer, whose import initializes CUDA.
+    """
+    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
+    if not cuda_home:
+        nvcc = shutil.which("nvcc")
+        cuda_home = (
+            os.path.dirname(os.path.dirname(nvcc)) if nvcc else _DEFAULT_CUDA_HOME
+        )
+    return shutil.which(os.path.join(cuda_home, "bin", "nvcc"))
+
 
 @functools.cache
 def has_flashinfer_cubin() -> bool:
@@ -99,12 +116,15 @@ def has_flashinfer() -> bool:
     if importlib.util.find_spec("flashinfer") is None:
         logger.debug_once("FlashInfer unavailable since package was not found")
         return False
-    # When not using flashinfer cubin,
-    # Also check if nvcc is available since it's required to JIT compile flashinfer
-    if not has_flashinfer_cubin() and shutil.which("nvcc") is None:
-        logger.debug_once(
-            "FlashInfer unavailable since nvcc was not found "
-            "and not using pre-downloaded cubins"
+    # FlashInfer's JIT runs nvcc and `ninja` (from PATH).
+    if not has_flashinfer_cubin() and (
+        _flashinfer_nvcc_path() is None or shutil.which("ninja") is None
+    ):
+        logger.warning_once(
+            "FlashInfer kernels are disabled: flashinfer-cubin is not installed "
+            "and nvcc (CUDA_HOME, CUDA_PATH, PATH or /usr/local/cuda) or ninja "
+            "(PATH) is missing. Set CUDA_HOME to a CUDA toolkit and put ninja on "
+            "PATH, or run `flashinfer download-kernels`."
         )
         return False
     return True