[Frontend][Build] Add vllm download-kernels for FlashInfer precompiled kernels (#58765)

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>

.buildkite/test_areas/entrypoints.yaml

@@ -31,6 +31,7 @@ steps:
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/entrypoints
+  - vllm/utils/flashinfer.py
   - tests/entrypoints/unit_tests
   - tests/entrypoints/launchers
   commands:

README.md

@@ -67,6 +67,7 @@ Install vLLM with [`uv`](https://docs.astral.sh/uv/) (recommended) or `pip`:
 
 ```bash
 uv pip install vllm
+vllm download-kernels  # precompiled kernels for faster startup
 ```
 
 Or [build from source](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source) for development.

docker/Dockerfile

@@ -936,13 +936,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \
     fi && \
     rm /tmp/requirements-cuda.txt /tmp/common.txt
 
-# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
-# https://docs.flashinfer.ai/installation.html
-# From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.7.0.post1
-RUN --mount=type=cache,target=/opt/uv/cache \
-    uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
-        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+# Install FlashInfer's precompiled kernels for all GPU architectures (the default).
+# `vllm download-kernels` wraps this command but vLLM is only installed in later
+# stages, and running it there would rebuild this large layer for every wheel.
+RUN --mount=type=cache,target=/opt/pip/cache \
+    PIP_CACHE_DIR=/opt/pip/cache flashinfer download-kernels
 
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES

docker/versions.json

@@ -73,9 +73,6 @@
     "RUN_WHEEL_CHECK": {
       "default": "true"
     },
-    "FLASHINFER_VERSION": {
-      "default": "0.7.0.post1"
-    },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"
     },

docs/cli/README.md

@@ -9,7 +9,7 @@ vllm --help
 Available Commands:
 
 ```bash
-vllm {chat,complete,serve,launch,bench,collect-env,run-batch,preload}
+vllm {chat,complete,serve,launch,bench,collect-env,download-kernels,run-batch,preload}
 ```
 
 ## serve
@@ -188,6 +188,21 @@ Start collecting environment information.
 vllm collect-env
 ```
 
+## download-kernels
+
+Install the precompiled kernels that vLLM's dependencies would otherwise download
+or compile at startup. Run it after installing or upgrading vLLM. It currently
+installs FlashInfer's kernels and does nothing when FlashInfer is not installed.
+
+```bash
+vllm download-kernels
+
+# Print the install commands without running them
+vllm download-kernels --dry-run
+```
+
+See [vllm download-kernels](./download-kernels.md) for the full reference of all available arguments.
+
 ## run-batch
 
 Run batch prompts and write results to file.

docs/contributing/README.md

@@ -41,6 +41,7 @@ If you are only developing vLLM's Python code, install vLLM using:
 
 ```bash
 VLLM_USE_PRECOMPILED=1 uv pip install -e .
+vllm download-kernels
 ```
 
 To rebuild only the Rust frontend binary:

docs/getting_started/installation/gpu.cuda.inc.md

@@ -23,14 +23,28 @@ Therefore, it is recommended to install vLLM with a **fresh new** environment. I
 
 ```bash
 uv pip install vllm --torch-backend=auto
+vllm download-kernels
 ```
 
 ??? console "pip"
     ```bash
     # Install vLLM with CUDA 12.9.
     pip install vllm --extra-index-url https://download.pytorch.org/whl/cu129
+    vllm download-kernels
     ```
 
+`vllm download-kernels` installs [FlashInfer](https://docs.flashinfer.ai/)'s
+precompiled kernels for the installed FlashInfer and CUDA versions. Without them,
+vLLM downloads and compiles kernels at startup, which can add several minutes on
+Hopper and newer GPUs, **especially Blackwell**, and logs a warning. Run it in the
+same Python environment after every vLLM install or upgrade, and in custom
+container images. The kernels are tied to the FlashInfer version, and FlashInfer
+refuses to import when they do not match, so an upgrade needs them refreshed.
+With a PyTorch built for CUDA older than 12.9, only `flashinfer-cubin` is
+installed, because FlashInfer does not publish `flashinfer-jit-cache` for it.
+For CUDA overrides and nightly kernels, use the
+[FlashInfer CLI](https://docs.flashinfer.ai/cli.html#download-kernels) directly.
+
 We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu130`), set `--torch-backend=cu130` (or `UV_TORCH_BACKEND=cu130`). If this doesn't work, try running `uv self update` to update `uv` first.
 
 !!! note
@@ -94,9 +108,10 @@ If you only need to change Python code, you can build and install vLLM without c
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
 VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto
+vllm download-kernels
 ```
 
-This command will do the following:
+The install command will do the following:
 
 1. Look for the current branch in your vLLM clone.
 1. Identify the corresponding base commit in the main branch.

docs/getting_started/quickstart.md

@@ -25,8 +25,11 @@ This guide will help you quickly get started with vLLM to perform:
     uv venv --python 3.12 --seed
     source .venv/bin/activate
     uv pip install vllm --torch-backend=auto
+    vllm download-kernels
     ```
 
+    `vllm download-kernels` installs precompiled kernels that vLLM otherwise downloads and compiles at startup, which can add several minutes on Hopper and newer GPUs. See the [installation guide](installation/gpu.md) for details.
+
     `uv` can [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu126`), set `--torch-backend=cu126` (or `UV_TORCH_BACKEND=cu126`).
 
     Another delightful way is to use `uv run` with `--with [dependency]` option, which allows you to run commands such as `vllm serve` without creating any permanent environment:
@@ -329,5 +332,5 @@ Some of the available backend options include:
 - On AMD ROCm: `TRITON_ATTN`, `ROCM_ATTN`, `ROCM_AITER_FA`, `ROCM_AITER_UNIFIED_ATTN`, `TRITON_MLA`, `ROCM_AITER_MLA` or `ROCM_AITER_TRITON_MLA`.
 - On Intel XPU: `FLASH_ATTN`, `TRITON_ATTN`, `TRITON_MLA`, `XPU_MLA_SPARSE`, `TORCH_SDPA` or `TURBOQUANT`.
 
-!!! warning
-    There are no pre-built vllm wheels containing Flash Infer, so you must install it in your environment first. Refer to the [Flash Infer official docs](https://docs.flashinfer.ai/) or see [docker/Dockerfile](../../docker/Dockerfile) for instructions on how to install it.
+!!! tip
+    FlashInfer is installed with vLLM, but its precompiled kernels are not. Run `vllm download-kernels` after installing or upgrading vLLM so they are not downloaded and compiled at startup.

docs/mkdocs/gen_files/generate_argparse.py

@@ -178,6 +178,9 @@ LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcomman
 launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION")
 RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand")
 PreloadSubcommand = auto_mock("vllm.entrypoints.cli.preload", "PreloadSubcommand")
+DownloadKernelsSubcommand = auto_mock(
+    "vllm.entrypoints.cli.download_kernels", "DownloadKernelsSubcommand"
+)
 sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS")
 openai_cli_args = auto_mock("vllm.entrypoints.launchers", "cli_args")
 run_batch = auto_mock("vllm.entrypoints.launchers", "run_batch")
@@ -332,6 +335,10 @@ pages = {
     "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False),
     "cli/run-batch.md": (create_parser(run_batch.make_arg_parser), True),
     "cli/preload.md": (create_parser(PreloadSubcommand.add_cli_args), True),
+    "cli/download-kernels.md": (
+        create_parser(DownloadKernelsSubcommand.add_cli_args),
+        False,
+    ),
     "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True),
     "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True),
     # URL kept as `mm_processor` for back-compat; command name is `mm-processor`

requirements/cuda.txt

@@ -10,12 +10,7 @@ torchaudio==2.11.0
 torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 torchcodec >= 0.14
 PyNvVideoCodec==2.0.4
-# FlashInfer should be updated together with the Dockerfile
-# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
-# install_requires so the published wheel does not carry an unresolvable pin
---extra-index-url https://flashinfer.ai/whl/
 flashinfer-python==0.7.0.post1
-flashinfer-cubin==0.7.0.post1
 apache-tvm-ffi==0.1.11
 tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1

tests/entrypoints/unit_tests/test_download_kernels_cli.py

@@ -0,0 +1,113 @@
+# 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

vllm/entrypoints/cli/download_kernels.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The `download-kernels` subcommand for the vLLM CLI.
+
+Installs the precompiled kernels that vLLM's dependencies would otherwise
+download or compile at startup. Run it after installing or upgrading vLLM:
+
+    vllm download-kernels
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+import typing
+from importlib.util import find_spec
+
+from vllm.entrypoints.cli.types import CLISubcommand
+from vllm.logger import init_logger
+
+if typing.TYPE_CHECKING:
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+else:
+    FlexibleArgumentParser = argparse.ArgumentParser
+
+logger = init_logger(__name__)
+
+
+def _download_flashinfer_kernels(dry_run: bool) -> int:
+    if find_spec("flashinfer") is None:
+        logger.info("FlashInfer is not installed; skipping its kernels.")
+        return 0
+    from vllm.utils.flashinfer import has_flashinfer_jit_cache_wheels
+
+    cmd = [sys.executable, "-m", "flashinfer"]
+    if has_flashinfer_jit_cache_wheels():
+        cmd.append("download-kernels")
+    else:
+        logger.info(
+            "FlashInfer does not publish flashinfer-jit-cache for this CUDA "
+            "version; installing flashinfer-cubin only."
+        )
+        cmd.append("install-cubin-wheel")
+    if dry_run:
+        cmd.append("--dry-run")
+    # FlashInfer refuses to import while kernels from another FlashInfer version
+    # are installed, which would stop its own CLI from replacing them.
+    env = {**os.environ, "FLASHINFER_DISABLE_VERSION_CHECK": "1"}
+    return subprocess.run(cmd, env=env, check=False).returncode
+
+
+class DownloadKernelsSubcommand(CLISubcommand):
+    """The `download-kernels` subcommand for the vLLM CLI."""
+
+    name = "download-kernels"
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+        parser.add_argument(
+            "--dry-run",
+            action="store_true",
+            help="Print the install commands without running them.",
+        )
+        return parser
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        sys.exit(_download_flashinfer_kernels(args.dry_run))
+
+    def subparser_init(
+        self, subparsers: argparse._SubParsersAction
+    ) -> FlexibleArgumentParser:
+        parser = subparsers.add_parser(
+            self.name,
+            help="Install precompiled kernels for faster startup.",
+            description=(
+                "Install the precompiled kernels that vLLM's dependencies would "
+                "otherwise download or compile at startup. Run it after "
+                "installing or upgrading vLLM."
+            ),
+            usage="vllm download-kernels [--dry-run]",
+        )
+        return DownloadKernelsSubcommand.add_cli_args(parser)
+
+
+def cmd_init() -> list[CLISubcommand]:
+    return [DownloadKernelsSubcommand()]

vllm/entrypoints/cli/main.py

@@ -42,6 +42,7 @@ def main():
 
     import vllm.entrypoints.cli.benchmark.main
     import vllm.entrypoints.cli.collect_env
+    import vllm.entrypoints.cli.download_kernels
     import vllm.entrypoints.cli.launch
     import vllm.entrypoints.cli.openai
     import vllm.entrypoints.cli.preload
@@ -60,6 +61,7 @@ def main():
         vllm.entrypoints.cli.launch,
         vllm.entrypoints.cli.benchmark.main,
         vllm.entrypoints.cli.collect_env,
+        vllm.entrypoints.cli.download_kernels,
         vllm.entrypoints.cli.preload,
         vllm.entrypoints.cli.run_batch,
         vllm.entrypoints.cli.snapshot,

vllm/utils/flashinfer.py

@@ -8,6 +8,7 @@ Users of vLLM should always import **only** these wrappers.
 import contextlib
 import functools
 import importlib
+import importlib.metadata
 import importlib.util
 import os
 import shutil
@@ -16,6 +17,7 @@ from typing import Any, NoReturn
 
 import requests
 import torch
+from packaging.version import Version
 
 import vllm.envs as envs
 from vllm.logger import init_logger
@@ -98,12 +100,60 @@ def has_flashinfer() -> bool:
             "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`."
+            "PATH, or run `vllm download-kernels`."
         )
         return False
     return True
 
 
+def _installed_version(distribution: str) -> str | None:
+    try:
+        return Version(importlib.metadata.version(distribution)).public
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def has_flashinfer_jit_cache_wheels() -> bool:
+    """Return whether FlashInfer publishes flashinfer-jit-cache for the installed
+    PyTorch's CUDA version, which it does from CUDA 12.9."""
+    cuda_version = torch.version.cuda
+    return cuda_version is not None and Version(cuda_version) >= Version("12.9")
+
+
+def warn_if_flashinfer_kernels_missing() -> None:
+    """Warn on Hopper and newer GPUs when FlashInfer's precompiled kernels are
+    missing or were installed for another FlashInfer version."""
+    if not (
+        current_platform.is_cuda()
+        and current_platform.has_device_capability(90)
+        and has_flashinfer()
+    ):
+        return
+    flashinfer_version = _installed_version("flashinfer-python")
+    packages = []
+    if not envs.VLLM_HAS_FLASHINFER_CUBIN:
+        packages.append("flashinfer-cubin")
+    if has_flashinfer_jit_cache_wheels():
+        packages.append("flashinfer-jit-cache")
+    installed = {name: _installed_version(name) for name in packages}
+    if flashinfer_version is None or all(
+        version == flashinfer_version for version in installed.values()
+    ):
+        return
+    logger.warning_once(
+        "FlashInfer's precompiled kernels are missing or do not match "
+        "flashinfer-python %s (%s). FlashInfer then downloads and compiles "
+        "kernels at startup, which can take several minutes, or fails to import "
+        "mismatched ones. Run `vllm download-kernels` in this Python environment "
+        "to install them.",
+        flashinfer_version,
+        ", ".join(
+            f"{name} {version or 'not installed'}"
+            for name, version in installed.items()
+        ),
+    )
+
+
 @functools.cache
 def has_flashinfer_bf16_gemm() -> bool:
     """Return whether FlashInfer exposes the BF16 dense GEMM API."""
@@ -1243,4 +1293,6 @@ __all__ = [
     "should_use_flashinfer_for_blockscale_fp8_gemm",
     "is_flashinfer_fp8_blockscale_gemm_supported",
     "is_flashinfer_cudnn_fp8_prefill_attn_supported",
+    "has_flashinfer_jit_cache_wheels",
+    "warn_if_flashinfer_kernels_missing",
 ]

vllm/v1/worker/gpu_worker.py

@@ -73,6 +73,7 @@ from vllm.profiler.wrapper import (
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
 from vllm.tracing import instrument
+from vllm.utils.flashinfer import warn_if_flashinfer_kernels_missing
 from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
 from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
 from vllm.utils.mem_constants import GiB_bytes
@@ -497,6 +498,7 @@ class Worker(WorkerBase):
 
             if self.use_v2_model_runner:
                 logger.info_once("Using V2 Model Runner")
+            warn_if_flashinfer_kernels_missing()
 
             # Set random seed.
             set_random_seed(self.model_config.seed)