Misha GoinandClaude Opus 5.5
29e4d2d42e
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import importlib.metadata
|
|
import os
|
|
import subprocess
|
|
from importlib.util import find_spec
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import regex as re
|
|
|
|
import vllm.utils.flashinfer as fi
|
|
from vllm.entrypoints.cli import download_kernels
|
|
|
|
|
|
@pytest.mark.skipif(find_spec("flashinfer") is None, reason="requires FlashInfer")
|
|
def test_download_kernels_resolves_kernel_warning(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
capfd: pytest.CaptureFixture[str],
|
|
caplog_vllm: pytest.LogCaptureFixture,
|
|
disable_log_dedup,
|
|
):
|
|
"""Runs the installed FlashInfer CLI, so a FlashInfer release that renames
|
|
its kernel packages, CLI or version check fails here instead of leaving
|
|
users with a warning that its own fix cannot clear."""
|
|
# A vLLM upgrade leaves kernels from the previous FlashInfer behind, which
|
|
# FlashInfer rejects on import. The command must work in that state.
|
|
stale_cubin = tmp_path / "flashinfer_cubin"
|
|
stale_cubin.mkdir()
|
|
(stale_cubin / "__init__.py").write_text(
|
|
f'__version__ = "0"\ndef get_cubin_dir(): return "{stale_cubin}"\n'
|
|
)
|
|
monkeypatch.setenv(
|
|
"PYTHONPATH", os.pathsep.join([str(tmp_path), os.getenv("PYTHONPATH", "")])
|
|
)
|
|
|
|
assert download_kernels._download_flashinfer_kernels(dry_run=True) == 0
|
|
kernels = dict(re.findall(r"Requirement: ([\w-]+)==(\S+)", capfd.readouterr().out))
|
|
assert kernels
|
|
|
|
installed = {"flashinfer-python": importlib.metadata.version("flashinfer-python")}
|
|
|
|
def version(name: str) -> str:
|
|
if name not in installed:
|
|
raise importlib.metadata.PackageNotFoundError(name)
|
|
return installed[name]
|
|
|
|
def warns() -> bool:
|
|
caplog_vllm.clear()
|
|
fi.warn_if_flashinfer_kernels_missing()
|
|
return "Run `vllm download-kernels`" in caplog_vllm.text
|
|
|
|
monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False)
|
|
monkeypatch.setattr(importlib.metadata, "version", version)
|
|
monkeypatch.setattr(fi, "has_flashinfer", lambda: True)
|
|
monkeypatch.setattr(
|
|
fi,
|
|
"current_platform",
|
|
SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True),
|
|
)
|
|
|
|
assert warns()
|
|
installed.update(dict.fromkeys(kernels, "0"))
|
|
assert warns()
|
|
installed.update(kernels)
|
|
assert not warns()
|
|
|
|
|
|
def test_cuda_without_jit_cache_wheels_needs_only_cubin(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
caplog_vllm: pytest.LogCaptureFixture,
|
|
disable_log_dedup,
|
|
):
|
|
"""FlashInfer publishes no flashinfer-jit-cache below CUDA 12.9, so the
|
|
command must not fail on it and the warning must not ask for it."""
|
|
commands = []
|
|
installed = {"flashinfer-python": "1.0", "flashinfer-cubin": "1.0"}
|
|
|
|
def version(name: str) -> str:
|
|
if name not in installed:
|
|
raise importlib.metadata.PackageNotFoundError(name)
|
|
return installed[name]
|
|
|
|
def run(cmd: list[str], **kwargs) -> SimpleNamespace:
|
|
commands.append(cmd)
|
|
return SimpleNamespace(returncode=0)
|
|
|
|
monkeypatch.delenv("VLLM_HAS_FLASHINFER_CUBIN", raising=False)
|
|
monkeypatch.setattr(fi.torch.version, "cuda", "12.8")
|
|
monkeypatch.setattr(download_kernels, "find_spec", lambda name: object())
|
|
monkeypatch.setattr(subprocess, "run", run)
|
|
monkeypatch.setattr(importlib.metadata, "version", version)
|
|
monkeypatch.setattr(fi, "has_flashinfer", lambda: True)
|
|
monkeypatch.setattr(
|
|
fi,
|
|
"current_platform",
|
|
SimpleNamespace(is_cuda=lambda: True, has_device_capability=lambda _: True),
|
|
)
|
|
|
|
assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0
|
|
assert commands[0][-1] == "install-cubin-wheel"
|
|
fi.warn_if_flashinfer_kernels_missing()
|
|
assert "vllm download-kernels" not in caplog_vllm.text
|
|
|
|
|
|
def test_download_kernels_skips_without_flashinfer(monkeypatch: pytest.MonkeyPatch):
|
|
"""Builds without FlashInfer (CPU, ROCm, XPU) have nothing to download."""
|
|
monkeypatch.setattr(download_kernels, "find_spec", lambda name: None)
|
|
monkeypatch.setattr(subprocess, "run", pytest.fail)
|
|
|
|
assert download_kernels._download_flashinfer_kernels(dry_run=False) == 0
|