import ctypes
import importlib.util
import json
import logging
import os
import re
import shutil
import subprocess
import sys
import sysconfig
from pathlib import Path
from shutil import which
import torch
from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_rust.build import build_rust
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
def load_module_from_path(module_name, path):
spec = importlib.util.spec_from_file_location(module_name, path)
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
return module
ROOT_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)
PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX = re.compile(r"vllm/_rust_[^/]*\.so$")
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
rust_build = load_module_from_path(
"rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")
)
VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED
USE_PRECOMPILED_RUST_FRONTEND = (
envs.VLLM_USE_PRECOMPILED or envs.VLLM_USE_PRECOMPILED_RUST
)
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
def get_precompiled_rust_extension_paths() -> list[Path]:
return sorted((ROOT_DIR / "vllm").glob("_rust_*.so"))
def get_missing_precompiled_rust_extension_modules() -> list[str]:
present = {
path.name.split(".", 1)[0] for path in get_precompiled_rust_extension_paths()
}
return [
module_name
for module_name in rust_build.rust_py_extension_module_names()
if module_name not in present
]
def has_precompiled_rust_extensions() -> bool:
return not get_missing_precompiled_rust_extension_modules()
def is_metadata_only_build() -> bool:
return bool({"egg_info", "dist_info"}.intersection(sys.argv[1:]))
if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
VLLM_TARGET_DEVICE = "cpu"
elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")):
logger.warning(
"vLLM only supports Linux platform (including WSL) and MacOS."
"Building on %s, "
"so vLLM may not be able to run correctly",
sys.platform,
)
VLLM_TARGET_DEVICE = "empty"
elif sys.platform.startswith("linux") and os.getenv("VLLM_TARGET_DEVICE") is None:
if torch.version.hip is not None:
VLLM_TARGET_DEVICE = "rocm"
logger.info("Auto-detected ROCm")
elif torch.version.xpu is not None:
VLLM_TARGET_DEVICE = "xpu"
logger.info("Auto-detected XPU")
elif torch.version.cuda is not None:
VLLM_TARGET_DEVICE = "cuda"
logger.info("Auto-detected CUDA")
else:
VLLM_TARGET_DEVICE = "cpu"
def is_sccache_available() -> bool:
return which("sccache") is not None and not bool(
int(os.getenv("VLLM_DISABLE_SCCACHE", "0"))
)
def is_ccache_available() -> bool:
return which("ccache") is not None
def is_ninja_available() -> bool:
return which("ninja") is not None
def is_freethreaded():
return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
def should_bundle_tcmalloc() -> bool:
import platform
return (
VLLM_TARGET_DEVICE == "cpu"
and sys.platform.startswith("linux")
and platform.machine() in ("aarch64", "x86_64", "s390x")
)
def find_tcmalloc() -> Path | None:
try:
output = subprocess.check_output(
["ldconfig", "-p"],
text=True,
stderr=subprocess.DEVNULL,
)
except Exception:
return None
for library_pattern in (
r"\blibtcmalloc_minimal\.so\.(\d+)\b",
r"\blibtcmalloc\.so\.(\d+)\b",
):
candidates: list[tuple[int, Path]] = []
for line in output.splitlines():
match = re.search(library_pattern, line)
if match is None or "=>" not in line:
continue
candidate = Path(line.split("=>")[1].strip())
if candidate.exists():
candidates.append((int(match.group(1)), candidate))
if candidates:
return max(candidates, key=lambda item: item[0])[1]
return None
def bundle_tcmalloc(build_lib: str) -> None:
tcmalloc_library = find_tcmalloc()
if tcmalloc_library is None:
logger.warning(
"Failed to locate tcmalloc. For best performance, "
"please install tcmalloc (e.g. `sudo apt-get "
"install -y --no-install-recommends libtcmalloc-minimal4`)"
)
return
bundle_dir = os.path.join(build_lib, "vllm", "libs")
os.makedirs(bundle_dir, exist_ok=True)
bundle_path = os.path.join(bundle_dir, tcmalloc_library.name)
shutil.copy2(tcmalloc_library, bundle_path)
logger.info("Bundled tcmalloc into wheel: %s", bundle_path)
class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)
class cmake_build_ext(build_ext):
did_config: dict[str, bool] = {}
def compute_num_jobs(self):
num_jobs = envs.MAX_JOBS
if num_jobs is not None:
num_jobs = int(num_jobs)
logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
else:
try:
num_jobs = len(os.sched_getaffinity(0))
except AttributeError:
num_jobs = os.cpu_count()
nvcc_threads = None
if _is_cuda() and CUDA_HOME is not None:
try:
nvcc_version = get_nvcc_cuda_version()
if nvcc_version >= Version("11.2"):
nvcc_threads = envs.NVCC_THREADS
if nvcc_threads is not None:
nvcc_threads = int(nvcc_threads)
logger.info(
"Using NVCC_THREADS=%d as the number of nvcc threads.",
nvcc_threads,
)
else:
nvcc_threads = 1
num_jobs = max(1, num_jobs // nvcc_threads)
except Exception as e:
logger.warning("Failed to get NVCC version: %s", e)
return num_jobs, nvcc_threads
def configure(self, ext: CMakeExtension) -> None:
if ext.cmake_lists_dir in cmake_build_ext.did_config:
return
cmake_build_ext.did_config[ext.cmake_lists_dir] = True
default_cfg = "Debug" if self.debug else "RelWithDebInfo"
cfg = envs.CMAKE_BUILD_TYPE or default_cfg
cmake_args = [
"-DCMAKE_BUILD_TYPE={}".format(cfg),
"-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
]
verbose = envs.VERBOSE
if verbose:
cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"]
if is_sccache_available():
cmake_args += [
"-DCMAKE_C_COMPILER_LAUNCHER=sccache",
"-DCMAKE_CXX_COMPILER_LAUNCHER=sccache",
"-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache",
"-DCMAKE_HIP_COMPILER_LAUNCHER=sccache",
]
elif is_ccache_available():
cmake_args += [
"-DCMAKE_C_COMPILER_LAUNCHER=ccache",
"-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
"-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache",
"-DCMAKE_HIP_COMPILER_LAUNCHER=ccache",
]
cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)]
cmake_args += ["-DVLLM_PYTHON_PATH={}".format(":".join(sys.path))]
fc_base_dir = os.path.join(ROOT_DIR, ".deps")
fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir)
cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)]
num_jobs, nvcc_threads = self.compute_num_jobs()
if nvcc_threads:
cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)]
if is_ninja_available():
build_tool = ["-G", "Ninja"]
cmake_args += [
"-DCMAKE_JOB_POOL_COMPILE:STRING=compile",
"-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs),
]
else:
build_tool = []
if _is_cuda() and CUDA_HOME is not None:
cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"]
elif _is_hip() and ROCM_HOME is not None:
cmake_args += [f"-DROCM_PATH={ROCM_HOME}"]
other_cmake_args = os.environ.get("CMAKE_ARGS")
if other_cmake_args:
cmake_args += other_cmake_args.split()
subprocess.check_call(
["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args],
cwd=self.build_temp,
)
def build_extensions(self) -> None:
try:
subprocess.check_output(["cmake", "--version"])
except OSError as e:
raise RuntimeError("Cannot find CMake executable") from e
if not os.path.exists(self.build_temp):
os.makedirs(self.build_temp)
targets = []
def target_name(s: str) -> str:
return s.removeprefix("vllm.").removeprefix("vllm_flash_attn.")
for ext in self.extensions:
self.configure(ext)
targets.append(target_name(ext.name))
num_jobs, _ = self.compute_num_jobs()
build_args = [
"--build",
".",
f"-j={num_jobs}",
*[f"--target={name}" for name in targets],
]
subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)
for ext in self.extensions:
outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute()
if outdir == self.build_temp:
continue
prefix = outdir
for _ in range(ext.name.count(".")):
prefix = prefix.parent
install_args = [
"cmake",
"--install",
".",
"--prefix",
prefix,
"--component",
target_name(ext.name),
]
subprocess.check_call(install_args, cwd=self.build_temp)
def run(self):
super().run()
if should_bundle_tcmalloc():
bundle_tcmalloc(self.build_lib)
import glob
files = glob.glob(
os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "**", "*.py"),
recursive=True,
)
for file in files:
dst_file = os.path.join(
"vllm/vllm_flash_attn", file.split("vllm/vllm_flash_attn/")[-1]
)
print(f"Copying {file} to {dst_file}")
os.makedirs(os.path.dirname(dst_file), exist_ok=True)
self.copy_file(file, dst_file)
if _is_cuda() or _is_hip():
print(
f"Copying {self.build_lib}/vllm/third_party/triton_kernels "
"to vllm/third_party/triton_kernels"
)
shutil.copytree(
f"{self.build_lib}/vllm/third_party/triton_kernels",
"vllm/third_party/triton_kernels",
dirs_exist_ok=True,
)
if _is_cuda():
deep_gemm_build = os.path.join(
self.build_lib, "vllm", "third_party", "deep_gemm"
)
if os.path.exists(deep_gemm_build):
print(f"Copying {deep_gemm_build} to vllm/third_party/deep_gemm")
shutil.copytree(
deep_gemm_build,
"vllm/third_party/deep_gemm",
dirs_exist_ok=True,
)
fmha_sm100_build = os.path.join(
self.build_lib, "vllm", "third_party", "fmha_sm100"
)
if os.path.exists(fmha_sm100_build):
print(f"Copying {fmha_sm100_build} to vllm/third_party/fmha_sm100")
shutil.copytree(
fmha_sm100_build,
"vllm/third_party/fmha_sm100",
dirs_exist_ok=True,
)
tml_fa4_build = os.path.join(
self.build_lib, "vllm", "third_party", "tml_fa4"
)
if os.path.exists(tml_fa4_build):
print(f"Copying {tml_fa4_build} to vllm/third_party/tml_fa4")
shutil.copytree(
tml_fa4_build,
"vllm/third_party/tml_fa4",
dirs_exist_ok=True,
)
class precompiled_build_ext(build_ext):
"""Disables extension building when using precompiled binaries."""
def run(self) -> None:
return
def build_extensions(self) -> None:
print("Skipping build_ext: using precompiled extensions.")
return
class precompiled_build_rust(build_rust):
"""Skips local Rust builds when all precompiled Rust artifacts are present."""
def run(self) -> None:
missing = []
if not PRECOMPILED_RUST_FRONTEND_PATH.exists():
missing.append(str(PRECOMPILED_RUST_FRONTEND_PATH))
missing_rust_extensions = get_missing_precompiled_rust_extension_modules()
if missing_rust_extensions:
missing.extend(
str(ROOT_DIR / "vllm" / f"{module_name}*.so")
for module_name in missing_rust_extensions
)
if not missing:
logger.info(
"Skipping local Rust build: using precompiled %s and %s",
PRECOMPILED_RUST_FRONTEND_PATH,
get_precompiled_rust_extension_paths(),
)
return
logger.warning(
"Precompiled wheel did not provide all Rust artifacts (%s); "
"falling back to local Rust build.",
", ".join(missing),
)
super().run()
class precompiled_wheel_utils:
"""Extracts libraries and other files from an existing wheel."""
@staticmethod
def fetch_metadata_for_variant(
commit: str,
variant: str | None,
*,
rocm: bool = False,
) -> tuple[list[dict], str]:
"""
Fetches metadata for a specific variant of the precompiled wheel.
For non-ROCm, fetches vllm metadata.
For ROCm, discovers all first-level packages and combines their
metadata into a single list.
"""
import json
from html.parser import HTMLParser
from urllib.request import urlopen
variant_dir = f"{variant}/" if variant is not None else ""
if not rocm:
repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
meta_url = repo_url + "metadata.json"
print(f"Trying to fetch nightly build metadata from {meta_url}")
with urlopen(meta_url) as resp:
wheels = json.loads(resp.read().decode("utf-8"))
return wheels, repo_url
repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}"
class LinkParser(HTMLParser):
def __init__(self):
super().__init__()
self.links: list[str] = []
def handle_starttag(self, tag, attrs):
if tag == "a":
href = dict(attrs).get("href")
if href:
self.links.append(href)
with urlopen(repo_url) as resp:
parser = LinkParser()
parser.feed(resp.read().decode("utf-8"))
packages = [
href.rstrip("/")
for href in parser.links
if href.endswith("/")
and href not in ("../", "./", "/")
and "/" not in href.rstrip("/")
]
wheels: list[dict] = []
for package in packages:
meta_url = f"{repo_url}{package}/metadata.json"
print(f"Trying to fetch nightly build metadata from {meta_url}")
with urlopen(meta_url) as resp:
package_wheels = json.loads(resp.read().decode("utf-8"))
wheels.extend(package_wheels)
return wheels, repo_url
@staticmethod
def is_rocm_system() -> bool:
"""Detect ROCm without relying on torch (for build environment)."""
if os.getenv("ROCM_PATH"):
return True
if os.path.isdir("/opt/rocm"):
return True
if which("rocminfo") is not None:
return True
try:
import torch
return torch.version.hip is not None
except ImportError:
return False
@staticmethod
def detect_system_cuda_variant() -> str:
"""Auto-detect CUDA variant from torch, nvidia-smi, or env default."""
supported = {12: "cu129", 13: "cu130"}
if envs.is_set("VLLM_MAIN_CUDA_VERSION"):
v = envs.VLLM_MAIN_CUDA_VERSION
print(f"Using VLLM_MAIN_CUDA_VERSION={v}")
return "cu" + v.replace(".", "")[:3]
cuda_version = None
try:
import torch
cuda_version = torch.version.cuda
except Exception:
pass
if not cuda_version:
try:
out = subprocess.run(
["nvidia-smi"], capture_output=True, text=True, timeout=10
)
if m := re.search(r"CUDA Version:\s*(\d+\.\d+)", out.stdout):
cuda_version = m.group(1)
except Exception:
pass
if not cuda_version:
cuda_version = envs.VLLM_MAIN_CUDA_VERSION
major = int(cuda_version.split(".")[0])
variant = supported.get(major, supported[max(supported)])
print(f"Detected CUDA {cuda_version}, using variant {variant}")
return variant
@staticmethod
def rocm_version_to_variant(rocm_version: str) -> str:
"""Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723."""
return "rocm" + rocm_version.replace(".", "")
@staticmethod
def detect_system_rocm_variant() -> str | None:
"""Auto-detect the ROCm wheel variant from the installed ROCm stack."""
rocm_version = get_rocm_version()
if not rocm_version:
try:
import torch
rocm_version = torch.version.hip
except Exception:
pass
if not rocm_version:
return None
variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version)
print(f"Detected ROCm {rocm_version}, using variant {variant}")
return variant
@staticmethod
def fetch_available_rocm_variants(commit: str) -> list[str]:
"""List ROCm wheel variants published for a commit on wheels.vllm.ai."""
from urllib.request import urlopen
index_url = f"https://wheels.vllm.ai/rocm/{commit}/"
print(f"Fetching available ROCm variants from {index_url}")
try:
with urlopen(index_url) as resp:
html = resp.read().decode("utf-8")
except Exception as e:
logger.warning(
"Failed to fetch ROCm variant index for commit %s: %s", commit, e
)
return []
variants = sorted(set(re.findall(r"rocm\d+", html)))
print(f"Available ROCm variants for commit {commit}: {variants}")
return variants
@staticmethod
def resolve_rocm_wheel_variant(
commit: str, variant_override: str | None
) -> str | None:
"""Resolve a ROCm wheel variant for a commit from wheels.vllm.ai."""
env_variant = precompiled_wheel_utils.detect_system_rocm_variant()
available = precompiled_wheel_utils.fetch_available_rocm_variants(commit)
if variant_override is not None:
if env_variant and variant_override != env_variant:
logger.warning(
"VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the "
"detected environment ROCm variant %s; refusing to use a "
"different ROCm patch wheel",
variant_override,
env_variant,
)
return None
if available and variant_override not in available:
logger.warning(
"Requested ROCm variant %s is not available for commit %s "
"(available: %s)",
variant_override,
commit,
available,
)
return None
return variant_override
if env_variant is None:
if available:
print(
"Could not detect ROCm version from the environment; "
f"available variants for commit {commit}: {available}"
)
return None
if env_variant in available:
return env_variant
logger.warning(
"Environment ROCm variant %s is not available for commit %s "
"(available: %s). Precompiled wheels may not be compatible.",
env_variant,
commit,
available,
)
return None
@staticmethod
def warn_if_rocm_torch_version_mismatch(
wheels: list[dict], repo_url: str, arch: str
) -> None:
"""Warn if installed torch differs from the custom ROCm build on
wheels.vllm.ai and suggest the correct install command."""
try:
installed = torch.__version__
except Exception:
return
def _wheel_version(pkg: str) -> str | None:
for w in wheels:
if w.get("package_name") == pkg and arch in w.get("platform_tag", ""):
v = w["version"]
if w.get("variant") and "+" not in v:
v = f"{v}+{w['variant']}"
return v
return None
expected = _wheel_version("torch")
if expected is None or installed == expected:
return
pkgs = f"torch=={expected}"
triton_ver = _wheel_version("triton")
if triton_ver:
pkgs += f" triton=={triton_ver}"
logger.warning(
"Installed PyTorch %s does not match the custom build %s "
"shipped with vLLM ROCm wheels. The ABI may differ from "
"official releases. If you hit extension load errors, "
"reinstall from the vLLM index:\n"
" pip install %s --extra-index-url %s\n"
" uv pip install %s --extra-index-url %s "
"--index-strategy unsafe-best-match",
installed,
expected,
pkgs,
repo_url,
pkgs,
repo_url,
)
@staticmethod
def find_local_rocm_wheel() -> str | None:
"""Search for a local vllm wheel in common locations."""
import glob
for pattern in ["/vllm-workspace/dist/vllm-*.whl", "./dist/vllm-*.whl"]:
wheels = glob.glob(pattern)
if wheels:
return sorted(wheels)[-1]
return None
@staticmethod
def fetch_wheel_from_pypi_index(index_url: str, package: str = "vllm") -> str:
"""Fetch the latest wheel URL from a PyPI-style simple index."""
import platform
from html.parser import HTMLParser
from urllib.parse import urljoin
from urllib.request import urlopen
arch = platform.machine()
class WheelLinkParser(HTMLParser):
def __init__(self):
super().__init__()
self.wheels = []
def handle_starttag(self, tag, attrs):
if tag == "a":
for name, value in attrs:
if name == "href" and value.endswith(".whl"):
self.wheels.append(value)
simple_url = f"{index_url.rstrip('/')}/{package}/"
print(f"Fetching wheel list from {simple_url}")
with urlopen(simple_url) as resp:
html = resp.read().decode("utf-8")
parser = WheelLinkParser()
parser.feed(html)
for wheel in reversed(parser.wheels):
if arch in wheel:
if wheel.startswith("http"):
return wheel
return urljoin(simple_url, wheel)
raise ValueError(f"No compatible wheel found for {arch} at {simple_url}")
@staticmethod
def determine_wheel_url_rocm() -> tuple[str, str | None]:
"""Determine the precompiled wheel for ROCm."""
local_wheel = precompiled_wheel_utils.find_local_rocm_wheel()
if local_wheel is not None:
print(f"Found local ROCm wheel: {local_wheel}")
return local_wheel, None
import platform
arch = platform.machine()
commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
if not commit or len(commit) != 40:
print(
f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}"
", trying to fetch base commit in main branch"
)
commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
variant = precompiled_wheel_utils.resolve_rocm_wheel_variant(
commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
)
print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}")
wheels, repo_url = None, None
if variant is not None:
try:
wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
commit, variant, rocm=True
)
precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch(
wheels, repo_url, arch
)
except Exception as e:
logger.warning(
"Failed to fetch ROCm wheel metadata for variant %s: %s",
variant,
e,
)
if wheels is not None and repo_url is not None:
from urllib.parse import urljoin
for wheel in wheels:
if wheel.get("package_name") == "vllm" and arch in wheel.get(
"platform_tag", ""
):
print(f"Found precompiled wheel metadata: {wheel}")
if "path" not in wheel:
raise ValueError(f"Wheel metadata missing path: {wheel}")
wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"])
download_filename = wheel.get("filename")
print(f"Using precompiled wheel URL: {wheel_url}")
return wheel_url, download_filename
logger.warning(
"No precompiled vllm wheel found for architecture %s "
"from repo %s. All available wheels: %s",
arch,
repo_url,
wheels,
)
index_url = os.getenv(
"VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple"
)
print(f"Fetching ROCm precompiled wheel from {index_url}")
wheel_url = precompiled_wheel_utils.fetch_wheel_from_pypi_index(index_url)
download_filename = wheel_url.split("/")[-1].split("#")[0]
print(f"Using ROCm precompiled wheel: {wheel_url}")
return wheel_url, download_filename
@staticmethod
def determine_wheel_url() -> tuple[str, str | None]:
"""
Try to determine the precompiled wheel URL or path to use.
The order of preference is:
1. user-specified wheel location (can be either local or remote, via
VLLM_PRECOMPILED_WHEEL_LOCATION)
2. user-specified variant (VLLM_PRECOMPILED_WHEEL_VARIANT) from nightly repo
or auto-detected CUDA variant based on system (torch, nvidia-smi)
3. the default variant from nightly repo
If downloading from the nightly repo, the commit can be specified via
VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch
is used.
"""
wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
if wheel_location is not None:
print(f"Using user-specified precompiled wheel location: {wheel_location}")
return wheel_location, None
else:
if precompiled_wheel_utils.is_rocm_system():
return precompiled_wheel_utils.determine_wheel_url_rocm()
import platform
arch = platform.machine()
variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
if variant is None:
variant = precompiled_wheel_utils.detect_system_cuda_variant()
commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
if not commit or len(commit) != 40:
print(
f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}"
", trying to fetch base commit in main branch"
)
commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
print(f"Using precompiled wheel commit {commit} with variant {variant}")
try_default = False
wheels, repo_url, download_filename = None, None, None
try:
wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
commit, variant
)
except Exception as e:
logger.warning(
"Failed to fetch precompiled wheel metadata for variant %s: %s",
variant,
e,
)
try_default = True
if try_default:
print("Trying the default variant from remote")
wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
commit, None
)
assert wheels is not None and repo_url is not None, (
"Failed to fetch precompiled wheel metadata"
)
"""[{
"package_name": "vllm",
"version": "0.11.2.dev278+gdbc3d9991",
"build_tag": null,
"python_tag": "cp38",
"abi_tag": "abi3",
"platform_tag": "manylinux1_x86_64",
"variant": null,
"filename": "vllm-0.11.2.dev278+gdbc3d9991-cp38-abi3-manylinux1_x86_64.whl",
"path": "../vllm-0.11.2.dev278%2Bgdbc3d9991-cp38-abi3-manylinux1_x86_64.whl"
},
...]"""
from urllib.parse import urljoin
for wheel in wheels:
if wheel.get("package_name") == "vllm" and arch in wheel.get(
"platform_tag", ""
):
print(f"Found precompiled wheel metadata: {wheel}")
if "path" not in wheel:
raise ValueError(f"Wheel metadata missing path: {wheel}")
wheel_url = urljoin(repo_url, wheel["path"])
download_filename = wheel.get("filename")
print(f"Using precompiled wheel URL: {wheel_url}")
break
else:
raise ValueError(
f"No precompiled vllm wheel found for architecture {arch} "
f"from repo {repo_url}. All available wheels: {wheels}"
)
return wheel_url, download_filename
@staticmethod
def extract_precompiled_and_patch_package(
wheel_url_or_path: str,
download_filename: str | None,
*,
extract_extensions: bool,
extract_rust_frontend: bool,
) -> dict:
import tempfile
import zipfile
temp_dir = None
try:
if not os.path.isfile(wheel_url_or_path):
wheel_filename = download_filename or wheel_url_or_path.split("/")[-1]
temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
wheel_path = os.path.join(temp_dir, wheel_filename)
print(f"Downloading wheel from {wheel_url_or_path} to {wheel_path}")
from urllib.request import urlretrieve
urlretrieve(wheel_url_or_path, filename=wheel_path)
else:
wheel_path = wheel_url_or_path
print(f"Using existing wheel at {wheel_path}")
package_data_patch = {}
with zipfile.ZipFile(wheel_path) as wheel:
exact_members = set()
if extract_extensions:
exact_members.update(
{
"vllm/_C.abi3.so",
"vllm/_C_stable_libtorch.abi3.so",
"vllm/_moe_C_stable_libtorch.abi3.so",
"vllm/_qutlass_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_flashkda_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
"vllm/cumem_allocator.abi3.so",
"vllm/spinloop.abi3.so",
"vllm/fs_io_C.abi3.so",
"vllm/_rocm_C.abi3.so",
}
)
if extract_rust_frontend:
exact_members.add("vllm/vllm-rs")
flash_attn_regex = re.compile(
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
flash_attn_files_to_skip = {
"vllm/vllm_flash_attn/__init__.py",
"vllm/vllm_flash_attn/flash_attn_interface.py",
}
triton_kernels_regex = re.compile(
r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
flashmla_regex = re.compile(
r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*")
tml_fa4_regex = re.compile(r"vllm/third_party/tml_fa4/.*")
file_members = []
for member in wheel.filelist:
if member.filename in exact_members:
file_members.append(member)
continue
if (
extract_rust_frontend
and PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX.match(
member.filename
)
):
file_members.append(member)
continue
if not extract_extensions:
continue
if (
(
flash_attn_regex.match(member.filename)
and member.filename not in flash_attn_files_to_skip
)
or triton_kernels_regex.match(member.filename)
or flashmla_regex.match(member.filename)
or deep_gemm_regex.match(member.filename)
or tml_fa4_regex.match(member.filename)
or fmha_sm100_regex.match(member.filename)
):
file_members.append(member)
for file in file_members:
print(f"[extract] {file.filename}")
target_path = os.path.join(".", file.filename)
os.makedirs(os.path.dirname(target_path), exist_ok=True)
with (
wheel.open(file.filename) as src,
open(target_path, "wb") as dst,
):
shutil.copyfileobj(src, dst)
mode = file.external_attr >> 16
if mode:
os.chmod(target_path, mode)
pkg = os.path.dirname(file.filename).replace("/", ".")
package_data_patch.setdefault(pkg, []).append(
os.path.basename(file.filename)
)
return package_data_patch
finally:
if temp_dir is not None:
print(f"Removing temporary directory {temp_dir}")
shutil.rmtree(temp_dir)
@staticmethod
def get_base_commit_in_main_branch() -> str:
try:
curl_cmd = [
"curl",
"-s",
"https://api.github.com/repos/vllm-project/vllm/commits/main",
]
github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN"))
if github_token:
curl_cmd += [
"-H",
f"Authorization: token {github_token}",
]
resp_json = subprocess.check_output(curl_cmd).decode("utf-8")
upstream_main_commit = json.loads(resp_json)["sha"]
print(f"Upstream main branch latest commit: {upstream_main_commit}")
if envs.VLLM_DOCKER_BUILD_CONTEXT:
return upstream_main_commit
try:
subprocess.check_output(
["git", "cat-file", "-e", f"{upstream_main_commit}"]
)
except subprocess.CalledProcessError:
subprocess.check_call(
["git", "fetch", "https://github.com/vllm-project/vllm", "main"]
)
current_branch = (
subprocess.check_output(["git", "branch", "--show-current"])
.decode("utf-8")
.strip()
)
base_commit = (
subprocess.check_output(
["git", "merge-base", f"{upstream_main_commit}", current_branch]
)
.decode("utf-8")
.strip()
)
return base_commit
except ValueError as err:
raise ValueError(err) from None
except Exception as err:
logger.warning(
"Failed to get the base commit in the main branch. "
"Using the nightly wheel. The libraries in this "
"wheel may not be compatible with your dev branch: %s",
err,
)
return "nightly"
def _no_device() -> bool:
return VLLM_TARGET_DEVICE == "empty"
def _is_cuda() -> bool:
has_cuda = torch.version.cuda is not None
return VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu()
def _is_hip() -> bool:
return (
VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm"
) and torch.version.hip is not None
def _is_tpu() -> bool:
return VLLM_TARGET_DEVICE == "tpu"
def _is_cpu() -> bool:
return VLLM_TARGET_DEVICE == "cpu"
def _is_xpu() -> bool:
return VLLM_TARGET_DEVICE == "xpu"
def _build_custom_ops() -> bool:
return _is_cuda() or _is_hip()
def get_rocm_version():
try:
if ROCM_HOME is None:
return None
librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so"
if not librocm_core_file.is_file():
return None
librocm_core = ctypes.CDLL(librocm_core_file)
VerErrors = ctypes.c_uint32
get_rocm_core_version = librocm_core.getROCmVersion
get_rocm_core_version.restype = VerErrors
get_rocm_core_version.argtypes = [
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
ctypes.POINTER(ctypes.c_uint32),
]
major = ctypes.c_uint32()
minor = ctypes.c_uint32()
patch = ctypes.c_uint32()
if (
get_rocm_core_version(
ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
)
== 0
):
return f"{major.value}.{minor.value}.{patch.value}"
return None
except Exception:
return None
def get_nvcc_cuda_version() -> Version:
"""Get the CUDA version from nvcc.
Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
"""
assert CUDA_HOME is not None, "CUDA_HOME is not set"
nvcc_output = subprocess.check_output(
[CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True
)
output = nvcc_output.split()
release_idx = output.index("release") + 1
nvcc_cuda_version = parse(output[release_idx].split(",")[0])
return nvcc_cuda_version
def get_vllm_version() -> str:
if env_version := os.getenv("VLLM_VERSION_OVERRIDE"):
print(f"Overriding VLLM version with {env_version} from VLLM_VERSION_OVERRIDE")
os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version
return get_version(write_to="vllm/_version.py")
version = get_version(write_to="vllm/_version.py")
sep = "+" if "+" not in version else "."
if _no_device():
if envs.VLLM_TARGET_DEVICE == "empty":
version += f"{sep}empty"
elif _is_cuda():
if USE_PRECOMPILED_EXTENSIONS and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX:
version += f"{sep}precompiled"
else:
cuda_version = str(get_nvcc_cuda_version())
if cuda_version != envs.VLLM_MAIN_CUDA_VERSION:
cuda_version_str = cuda_version.replace(".", "")[:3]
if "sdist" not in sys.argv:
version += f"{sep}cu{cuda_version_str}"
elif _is_hip():
rocm_version = get_rocm_version() or torch.version.hip
if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION:
version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
elif _is_tpu():
version += f"{sep}tpu"
elif _is_cpu():
if VLLM_TARGET_DEVICE == "cpu":
version += f"{sep}cpu"
elif _is_xpu():
version += f"{sep}xpu"
else:
raise RuntimeError("Unknown runtime environment")
return version
def get_requirements() -> list[str]:
"""Get Python package dependencies from requirements.txt."""
requirements_dir = ROOT_DIR / "requirements"
def _read_requirements(filename: str) -> list[str]:
with open(requirements_dir / filename) as f:
requirements = f.read().strip().split("\n")
resolved_requirements = []
for line in requirements:
if line.startswith("-r "):
resolved_requirements += _read_requirements(line.split()[1])
elif (
not line.startswith("--")
and not line.startswith("#")
and line.strip() != ""
):
resolved_requirements.append(line)
return resolved_requirements
if _no_device():
requirements = _read_requirements("common.txt")
elif _is_cuda():
requirements = _read_requirements("cuda.txt")
cuda_major, cuda_minor = torch.version.cuda.split(".")
modified_requirements = []
for req in requirements:
if "vllm-flash-attn" in req and cuda_major != "12":
continue
if "flashinfer-cubin" in req:
continue
if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
if "humming-kernels[cu13]" in req and cuda_major == "12":
req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]")
modified_requirements.append(req)
requirements = modified_requirements
elif _is_hip():
requirements = _read_requirements("rocm.txt")
elif _is_tpu():
requirements = _read_requirements("tpu.txt")
elif _is_cpu():
requirements = _read_requirements("cpu.txt")
elif _is_xpu():
requirements = _read_requirements("xpu.txt")
else:
raise ValueError("Unsupported platform, please use CUDA, ROCm, or CPU.")
return requirements
ext_modules = []
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
if not _is_xpu() and sys.version_info >= (3, 11):
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
ext_modules.append(CMakeExtension(name="vllm.fs_io_C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
if USE_PRECOMPILED_EXTENSIONS or (
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3")
):
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
ext_modules.append(
CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True)
)
if USE_PRECOMPILED_EXTENSIONS or (
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
):
ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
ext_modules.append(
CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
)
if USE_PRECOMPILED_EXTENSIONS or (
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0")
):
ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True))
if envs.VLLM_USE_PRECOMPILED or (
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3")
):
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
ext_modules.append(CMakeExtension(name="vllm.tml_fa4", optional=True))
if _is_cpu():
import platform
if platform.machine() in ("x86_64", "AMD64"):
ext_modules.append(CMakeExtension(name="vllm._C"))
ext_modules.append(CMakeExtension(name="vllm._C_AVX512"))
ext_modules.append(CMakeExtension(name="vllm._C_AVX2"))
else:
ext_modules.append(CMakeExtension(name="vllm._C"))
if _build_custom_ops():
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
package_data = {
"vllm": [
"py.typed",
"libs/*.so*",
"model_executor/layers/fused_moe/configs/*.json",
"model_executor/layers/quantization/utils/configs/*.json",
"entrypoints/serve/instrumentator/static/*.js",
"entrypoints/serve/instrumentator/static/*.css",
"distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp",
"third_party/flash_linear_attention/LICENSE",
"third_party/deep_gemm/include/**/*.cuh",
"third_party/deep_gemm/include/**/*.h",
"third_party/deep_gemm/include/**/*.hpp",
"third_party/fmha_sm100/csrc/**/*.cu",
"third_party/fmha_sm100/csrc/**/*.h",
"third_party/fmha_sm100/csrc/**/*.jinja",
"third_party/fmha_sm100/csrc/**/*.cu.jinja",
"third_party/fmha_sm100/cute/**/*.cu",
"third_party/fmha_sm100/cutlass/include/**/*.h",
"third_party/fmha_sm100/cutlass/include/**/*.hpp",
"third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
"third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
]
}
def add_vllm_package_data(filename: str) -> None:
vllm_files = package_data.setdefault("vllm", [])
if filename not in vllm_files:
vllm_files.append(filename)
if USE_PRECOMPILED_RUST_FRONTEND and not is_metadata_only_build():
wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()
patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(
wheel_url,
download_filename,
extract_extensions=USE_PRECOMPILED_EXTENSIONS,
extract_rust_frontend=True,
)
for pkg, files in patch.items():
package_data.setdefault(pkg, []).extend(files)
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
add_vllm_package_data("vllm-rs")
for rust_extension_path in get_precompiled_rust_extension_paths():
add_vllm_package_data(rust_extension_path.name)
if _no_device():
ext_modules = []
if not ext_modules:
cmdclass = {}
else:
cmdclass = {
"build_ext": precompiled_build_ext
if USE_PRECOMPILED_EXTENSIONS
else cmake_build_ext,
}
if (
USE_PRECOMPILED_RUST_FRONTEND
or PRECOMPILED_RUST_FRONTEND_PATH.exists()
or has_precompiled_rust_extensions()
):
cmdclass["build_rust"] = precompiled_build_rust
vllm_version = get_vllm_version()
rust_build.prepare_build_environment()
rust_extensions = rust_build.rust_extensions(
optional=not should_require_rust_frontend()
)
setup(
version=vllm_version,
ext_modules=ext_modules,
rust_extensions=rust_extensions,
install_requires=get_requirements(),
extras_require={
"zen": ["zentorch==2.13.0.0"],
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
"instanttensor": ["instanttensor >= 0.1.9"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
"audio": [
"av",
"scipy",
"soundfile",
"soxr",
"mistral_common[audio]",
],
"video": [],
"deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"],
"flashinfer": [],
"b12x": ["b12x==1.2.4"],
"helion": ["helion==1.4.0"],
"grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"],
"otel": [
"opentelemetry-sdk>=1.26.0",
"opentelemetry-api>=1.26.0",
"opentelemetry-exporter-otlp>=1.26.0",
"opentelemetry-semantic-conventions-ai>=0.4.1",
],
"extra-quant": ["vllm-gguf-plugin>=0.0.2"],
},
cmdclass=cmdclass,
package_data=package_data,
)