diff --git a/tests/utils_/test_flashinfer_utils.py b/tests/utils_/test_flashinfer_utils.py new file mode 100644 index 0000000000..3d28486ff5 --- /dev/null +++ b/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() diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index abc4d75500..44da20a3ed 100644 --- a/vllm/utils/flashinfer.py +++ b/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