| 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() |