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>
88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
# 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()]
|