admin/1Cat-vLLM
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
main
1Cat-vLLM/setup.py
T
LeonandClaude Opus 5.5 7b1cbd1727 [Build][SM70] Fix the Docker source build and keep all SM70 extensions (#1207) 
Two gaps in the Docker source build (docker/Dockerfile --target build):

- setup.py preloads vllm/envs_metadata.py before vllm/envs.py and installs
  scripts/serve_flash_next_nvfp4_v100.sh, but the native csrc-build stage
  copies only envs.py, __init__.py and the 27B serve script, so the stage
  stops with FileNotFoundError before compiling anything. Copy both files.
- The build stage reinstalls vLLM from the native-stage wheel with
  VLLM_USE_PRECOMPILED=1, and setup.py copies only the native libraries its
  extraction patterns name. The SM70 pattern lists sampler, exact_reduce,
  sparse_attention and qsa_device, so _sm70_qsa_indexer_C (#1127) and
  _sm70_router_C (#1147) are built but left out of the final wheel. Without
  the indexer, the QSA shared-key path reports native_extension_unavailable
  and falls back to the paged Triton indexer; the router operator tests
  skip. Add both names.

The extraction test now derives the expected libraries from the SM70/H3
CMakeExtension declarations in setup.py, so a newly declared extension
cannot be dropped silently again.

Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-28 10:26:16 +08:00

1583 lines
60 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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.tags import sys_tags
from packaging.version import Version, parse
from setuptools import Extension, find_packages, setup
from setuptools.command.bdist_wheel import bdist_wheel
from setuptools.command.build_ext import build_ext
from setuptools_rust import Binding, RustExtension
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"
FLASH_ATTN_V100_ROOT = ROOT_DIR / "flash-attention-v100"
FLASH_ATTN_V100_PACKAGE = FLASH_ATTN_V100_ROOT / "flash_attn_v100"
FLASH_QLA_SM70_ROOT = (
ROOT_DIR / "flash_qla" / "ops" / "gated_delta_rule" / "chunk" / "sm70"
)
ONECAT_TORCH_CU128_URLS = {
"torch==2.10.0": ("torch", "2.10.0"),
"torchaudio==2.10.0": ("torchaudio", "2.10.0"),
"torchvision==0.25.0": ("torchvision", "0.25.0"),
}
def _onecat_torch_cu128_urls() -> dict[str, str]:
platform = sysconfig.get_platform()
if platform not in {"linux-x86_64", "linux-aarch64"}:
raise RuntimeError(
f"No pinned CUDA 12.8 PyTorch wheels for platform {platform}"
)
tag = next(sys_tags())
wheel_platform = "manylinux_2_28_" + platform.removeprefix("linux-")
return {
requirement: (
f"{name} @ https://download.pytorch.org/whl/cu128/"
f"{name}-{version}%2Bcu128-{tag.interpreter}-{tag.abi}-{wheel_platform}.whl"
)
for requirement, (name, version) in ONECAT_TORCH_CU128_URLS.items()
}
# cannot import envs directly because it depends on vllm,
# which is not installed yet
# Preload its dependency without importing vllm's runtime initialization.
load_module_from_path("vllm.envs_metadata", ROOT_DIR / "vllm" / "envs_metadata.py")
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED
# VLLM_USE_PRECOMPILED implies precompiled rust frontend too.
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")
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")
)
def find_tcmalloc() -> Path | None:
try:
# get all shared libs the dynamic loader knows about
output = subprocess.check_output(
["ldconfig", "-p"],
text=True,
stderr=subprocess.DEVNULL,
)
except Exception:
return None
# search for libtcmalloc and libtcmalloc_minimal
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:
# if multiple candidates are found, pick the one with the highest
# version number
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)
def bundle_flash_attn_v100(build_lib: str) -> None:
"""Build and embed the SM70 Flash-V100 package into the 1Cat wheel."""
if not FLASH_ATTN_V100_ROOT.exists():
logger.warning("Flash-V100 source tree is missing; skipping wheel bundle")
return
env = os.environ.copy()
env["TORCH_CUDA_ARCH_LIST"] = os.environ.get(
"FLASH_ATTN_V100_CUDA_ARCH_LIST", _volta_cuda_arch_list() or "7.0"
)
subprocess.check_call(
[sys.executable, "setup.py", "build_ext", "--inplace"],
cwd=FLASH_ATTN_V100_ROOT,
env=env,
)
dst_pkg = Path(build_lib) / "flash_attn_v100"
dst_pkg.mkdir(parents=True, exist_ok=True)
for py_file in FLASH_ATTN_V100_PACKAGE.glob("*.py"):
dst_file = dst_pkg / py_file.name
if py_file.resolve() != dst_file.resolve():
shutil.copy2(py_file, dst_file)
ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
for ext_name in ("flash_attn_v100_cuda", "paged_kv_utils"):
matches = sorted(FLASH_ATTN_V100_ROOT.glob(f"{ext_name}*{ext_suffix}"))
if not matches:
matches = sorted(FLASH_ATTN_V100_ROOT.glob(f"{ext_name}*.so"))
if not matches:
raise RuntimeError(f"Failed to build Flash-V100 extension: {ext_name}")
dst_ext = dst_pkg / matches[-1].name
shutil.copy2(matches[-1], dst_ext)
remove_rpath(dst_ext)
def bundle_precompiled_flash_attn_v100(build_lib: str) -> None:
"""Embed Flash-V100 files extracted from a precompiled 1Cat wheel."""
extracted_pkg = ROOT_DIR / "flash_attn_v100"
ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
extensions: list[Path] = []
for ext_name in ("flash_attn_v100_cuda", "paged_kv_utils"):
matches = sorted(extracted_pkg.glob(f"{ext_name}*{ext_suffix}"))
if not matches:
matches = sorted(extracted_pkg.glob(f"{ext_name}*.so"))
if not matches:
return
extensions.append(matches[-1])
dst_pkg = Path(build_lib) / "flash_attn_v100"
dst_pkg.mkdir(parents=True, exist_ok=True)
for py_file in FLASH_ATTN_V100_PACKAGE.glob("*.py"):
dst_file = dst_pkg / py_file.name
if py_file.resolve() != dst_file.resolve():
shutil.copy2(py_file, dst_file)
for extension in extensions:
dst_ext = dst_pkg / extension.name
shutil.copy2(extension, dst_ext)
remove_rpath(dst_ext)
def bundle_flash_qla_sm70(build_lib: str, build_temp: str) -> None:
"""Precompile the SM70 GDN extension so runtime never requires NVCC."""
if not _volta_cuda_arch_list():
return
src = FLASH_QLA_SM70_ROOT / "csrc" / "gdn_forward.cu"
if not src.exists():
raise RuntimeError(f"SM70 FlashQLA source is missing: {src}")
from torch.utils.cpp_extension import load as load_torch_extension
extension_build_dir = Path(build_temp) / "flash_qla_sm70_gdn_strided"
extension_build_dir.mkdir(parents=True, exist_ok=True)
previous_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
os.environ["TORCH_CUDA_ARCH_LIST"] = _volta_cuda_arch_list()
try:
extension = load_torch_extension(
name="flash_qla_sm70_gdn_strided",
sources=[str(src)],
build_directory=str(extension_build_dir),
extra_cuda_cflags=[
"-O3",
*torch.utils.cpp_extension._get_cuda_arch_flags(),
],
extra_cflags=["-O3"],
with_cuda=True,
verbose=bool(int(os.environ.get("FLASH_QLA_SM70_VERBOSE_BUILD", "0"))),
)
finally:
if previous_arch_list is None:
os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
else:
os.environ["TORCH_CUDA_ARCH_LIST"] = previous_arch_list
dst_pkg = (
Path(build_lib) / "flash_qla" / "ops" / "gated_delta_rule" / "chunk" / "sm70"
)
dst_pkg.mkdir(parents=True, exist_ok=True)
dst_ext = dst_pkg / Path(extension.__file__).name
shutil.copy2(extension.__file__, dst_ext)
remove_rpath(dst_ext)
logger.info("Bundled SM70 FlashQLA extension into wheel: %s", dst_ext)
def remove_rpath(path: Path) -> None:
patchelf = which("patchelf")
if patchelf is None:
raise RuntimeError(
"patchelf is required to build a portable 1Cat SM70 wheel; "
f"refusing to leave a build-host RPATH in {path}"
)
subprocess.run([patchelf, "--remove-rpath", str(path)], check=True)
class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
kwa.setdefault("py_limited_api", not is_freethreaded())
super().__init__(name, sources=[], **kwa)
self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)
class cmake_build_ext(build_ext):
# A dict of extension directories that have been configured.
did_config: dict[str, bool] = {}
#
# Determine number of compilation jobs and optionally nvcc compile threads.
#
def compute_num_jobs(self):
# `num_jobs` is either the value of the MAX_JOBS environment variable
# (if defined) or the number of CPUs available.
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:
# os.sched_getaffinity() isn't universally available, so fall
# back to os.cpu_count() if we get an error here.
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` is either the value of the NVCC_THREADS
# environment variable (if defined) or 1.
# when it is set, we reduce `num_jobs` to avoid
# overloading the system.
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
#
# Perform cmake configuration for a single extension.
#
def configure(self, ext: CMakeExtension) -> None:
# If we've already configured using the CMakeLists.txt for
# this extension, exit early.
if ext.cmake_lists_dir in cmake_build_ext.did_config:
return
cmake_build_ext.did_config[ext.cmake_lists_dir] = True
# Select the build type.
# Note: optimization level + debug info are set by the build type
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",
]
# Pass the python executable to cmake so it can find an exact
# match.
cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)]
# Pass the python path to cmake so it can reuse the build dependencies
# on subsequent calls to python.
cmake_args += ["-DVLLM_PYTHON_PATH={}".format(":".join(sys.path))]
# Override the base directory for FetchContent downloads to $ROOT/.deps
# This allows sharing dependencies between profiles,
# and plays more nicely with sccache.
# To override this, set the FETCHCONTENT_BASE_DIR environment variable.
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)]
#
# Setup parallelism and build tool
#
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:
# Default build tool to whatever cmake picks.
build_tool = []
# Make sure we use the nvcc from CUDA_HOME
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:
# Ensure that CMake is present and working
try:
subprocess.check_output(["cmake", "--version"])
except OSError as e:
raise RuntimeError("Cannot find CMake executable") from e
# Create build directory if it does not exist.
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.")
# Build all the extensions
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)
# Install the libraries
for ext in self.extensions:
# Install the extension into the proper location
outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute()
# Skip if the install directory is the same as the build directory
if outdir == self.build_temp:
continue
# CMake appends the extension prefix to the install path,
# and outdir already contains that prefix, so we need to remove it.
prefix = outdir
for _ in range(ext.name.count(".")):
prefix = prefix.parent
# prefix here should actually be the same for all components
install_args = [
"cmake",
"--install",
".",
"--prefix",
prefix,
"--component",
target_name(ext.name),
]
subprocess.check_call(install_args, cwd=self.build_temp)
def run(self):
# First, run the standard build_ext command to compile the extensions
super().run()
# bundle tcmalloc into CPU wheels for best OOB perf
if should_bundle_tcmalloc():
bundle_tcmalloc(self.build_lib)
if _is_cuda():
# In-place imports resolve the source packages, not build/lib.
# Bundle companion kernels beside them just like the native _C
# extensions; wheel builds still stage everything in build_lib.
bundle_dir = str(ROOT_DIR) if self.inplace else self.build_lib
bundle_flash_attn_v100(
str(FLASH_ATTN_V100_ROOT) if self.inplace else self.build_lib
)
bundle_flash_qla_sm70(bundle_dir, self.build_temp)
# copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current
# directory so that they can be included in the editable build
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():
# copy vllm/third_party/triton_kernels/**/*.py from self.build_lib
# to current directory so that they can be included in the editable
# build
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():
# copy vendored deep_gemm package from build_lib to source tree
# for editable installs
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,
)
class vllm_bdist_wheel(bdist_wheel):
"""Use interpreter-specific tags when native libraries require the full ABI."""
def finalize_options(self) -> None:
super().finalize_options()
full_api_extensions = any(
not extension.py_limited_api
for extension in self.distribution.ext_modules or []
)
# Precompiled wheels supply native libraries through package_data, even
# when TORCH_CUDA_ARCH_LIST does not declare their CMake targets.
cpython_package_data = any(
re.search(r"\.cpython-\d+[^/]*\.so$", filename)
for filenames in (self.distribution.package_data or {}).values()
for filename in filenames
)
if full_api_extensions or cpython_package_data:
self.root_is_pure = False
if self.py_limited_api:
logger.warning(
"Ignoring --py-limited-api=%s: wheel contains "
"CPython-specific native extensions",
self.py_limited_api,
)
self.py_limited_api = None
class precompiled_build_ext(build_ext):
"""Disables extension building when using precompiled binaries."""
def run(self) -> None:
if _is_cuda():
bundle_precompiled_flash_attn_v100(
str(FLASH_ATTN_V100_ROOT) if self.inplace else self.build_lib
)
def build_extensions(self) -> None:
print("Skipping build_ext: using precompiled extensions.")
return
class precompiled_build_rust(build_rust):
"""Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
def run(self) -> None:
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
logger.info(
"Skipping local Rust build: using precompiled %s",
PRECOMPILED_RUST_FRONTEND_PATH,
)
return
logger.warning(
"Precompiled wheel did not provide %s; falling back to local Rust build.",
PRECOMPILED_RUST_FRONTEND_PATH,
)
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
) -> tuple[list[dict], str]:
"""
Fetches metadata for a specific variant of the precompiled wheel.
"""
variant_dir = f"{variant}/" if variant is not None else ""
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}")
from urllib.request import urlopen
with urlopen(meta_url) as resp:
# urlopen raises HTTPError on unexpected status code
wheels = json.loads(resp.read().decode("utf-8"))
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."""
# Map CUDA major version to hosted wheel variants on wheels.vllm.ai
supported = {12: "cu129", 13: "cu130"}
# Respect explicitly set VLLM_MAIN_CUDA_VERSION
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]
# Try torch.version.cuda
cuda_version = None
try:
import torch
cuda_version = torch.version.cuda
except Exception:
pass
# Try nvidia-smi
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
# Fall back to default
if not cuda_version:
cuda_version = envs.VLLM_MAIN_CUDA_VERSION
# Map to supported variant
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 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."""
# Search for local wheel first
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
# Fall back to AMD's PyPI index
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:
# ROCm: use local wheel or AMD's PyPI index
# TODO: When we have ROCm nightly wheels, we can update this logic.
if precompiled_wheel_utils.is_rocm_system():
return precompiled_wheel_utils.determine_wheel_url_rocm()
import platform
arch = platform.machine()
# try to fetch the wheel metadata from the nightly wheel repo,
# detecting CUDA variant from system if not specified
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 # try outside handler to keep the stacktrace simple
if try_default:
print("Trying the default variant from remote")
wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
commit, None
)
# if this also fails, then we have nothing more to try / cache
assert wheels is not None and repo_url is not None, (
"Failed to fetch precompiled wheel metadata"
)
# The metadata.json has the following format:
# see .buildkite/scripts/generate-nightly-index.py for details
"""[{
"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:
# TODO: maybe check more compatibility later? (python_tag, abi_tag, etc)
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):
# use provided filename first, then derive from URL
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/_C_gguf.abi3.so",
"vllm/_moe_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa2_C_sm75.abi3.so",
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
"vllm/cumem_allocator.abi3.so",
"vllm/spinloop.abi3.so",
# ROCm-specific libraries
"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"
)
# __init__.py and flash_attn_interface.py are source-controlled
# in vllm and should not be overwritten (matches cmake exclusions)
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"
)
# DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.)
deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
flash_attn_v100_ext_regex = re.compile(
r"flash_attn_v100/"
r"(?:flash_attn_v100_cuda|paged_kv_utils).*\.so"
)
flash_qla_sm70_ext_regex = re.compile(
r"flash_qla/ops/gated_delta_rule/chunk/sm70/[^/]+\.so"
)
sm70_ext_regex = re.compile(
r"vllm/_sm70_(?:sampler|exact_reduce|sparse_attention|qsa_device"
r"|qsa_indexer|router)_C(?:\.[^/]+)?\.so$"
)
h3_ext_regex = re.compile(
r"vllm/_h3_(?:w8a16|flashinfer|flashattn)_C(?:\.[^/]+)?\.so$"
)
file_members = []
for member in wheel.filelist:
if member.filename in exact_members:
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 flash_attn_v100_ext_regex.match(member.filename)
or flash_qla_sm70_ext_regex.match(member.filename)
or sm70_ext_regex.match(member.filename)
or h3_ext_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)
if flash_attn_v100_ext_regex.match(file.filename):
remove_rpath(Path(target_path))
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:
# Get the latest commit hash of the upstream main branch.
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}")
# In Docker build context, .git may be immutable or missing.
if envs.VLLM_DOCKER_BUILD_CONTEXT:
return upstream_main_commit
# Check if the upstream_main_commit exists in the local repo
try:
subprocess.check_output(
["git", "cat-file", "-e", f"{upstream_main_commit}"]
)
except subprocess.CalledProcessError:
# If not present, fetch it from the remote repository.
# Note that this does not update any local branches,
# but ensures that this commit ref and its history are
# available in our local repo.
subprocess.check_call(
["git", "fetch", "https://github.com/vllm-project/vllm", "main"]
)
# Then get the commit hash of the current branch that is the same as
# the upstream main commit.
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():
# Get the Rocm version from the ROCM_HOME/bin/librocm-core.so
# see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21
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 _cuda_arches_from_env() -> list[tuple[int, int]] | None:
"""Return TORCH_CUDA_ARCH_LIST entries as (major, minor), if explicit."""
raw_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
if not raw_arch_list:
return None
arches: list[tuple[int, int]] = []
for raw_arch in re.split(r"[;,\s]+", raw_arch_list):
arch = raw_arch.strip().lower()
if not arch:
continue
arch = arch.replace("+ptx", "")
arch = arch.removeprefix("sm_").removeprefix("compute_")
match = re.fullmatch(r"(\d+)(?:\.(\d+))?", arch)
if match is None:
return None
major_str, minor_str = match.groups()
if minor_str is None and len(major_str) >= 2:
major = int(major_str[:-1])
minor = int(major_str[-1])
else:
major = int(major_str)
minor = int(minor_str or 0)
arches.append((major, minor))
return arches or None
def _cuda_arch_at_least(major: int, minor: int = 0) -> bool:
arches = _cuda_arches_from_env()
if arches is None:
return True
return any(arch >= (major, minor) for arch in arches)
def _cuda_arch_contains(major: int, minor: int = 0) -> bool:
arches = _cuda_arches_from_env()
return arches is not None and (major, minor) in arches
def _volta_cuda_arch_list() -> str:
"""Keep the requested Volta targets, including explicit PTX, for bundles."""
aliases = {"70": "7.0", "7.0": "7.0", "72": "7.2", "7.2": "7.2"}
arches = []
for raw in re.split(r"[;,\s]+", os.environ.get("TORCH_CUDA_ARCH_LIST", "")):
lowered = raw.strip().lower()
arch = lowered.removesuffix("+ptx")
arch = arch.removeprefix("sm_").removeprefix("compute_")
if arch not in aliases:
continue
target = aliases[arch] + ("+PTX" if lowered.endswith("+ptx") else "")
if target not in arches:
arches.append(target)
return ";".join(arches)
def get_vllm_version() -> str:
# Allow overriding the version. This is useful to build platform-specific
# wheels (e.g. CPU, TPU) without modifying the source.
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 "." # dev versions might contain +
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]
# skip this for source tarball, required for pypi
if "sdist" not in sys.argv:
version += f"{sep}cu{cuda_version_str}"
elif _is_hip():
# Get the Rocm Version
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():
# Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above),
# not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+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(".")
# An SM70 wheel is a CUDA 12.8 release artifact. Keep its runtime
# dependency on the matching PyTorch wheel in the wheel metadata so a
# user installing the wheel does not need to know about this build
# detail or export a local-environment override. The explicit env
# switch remains for reproducible packaging jobs that pin the same
# dependency while using a non-SM70 target.
pin_torch_cu128 = bool(int(os.getenv("ONECAT_VLLM_PIN_TORCH_CU128", "0"))) or (
(_cuda_arch_contains(7, 0) or _cuda_arch_contains(7, 2))
and torch.version.cuda == "12.8"
)
torch_cu128_urls = _onecat_torch_cu128_urls() if pin_torch_cu128 else {}
if pin_torch_cu128:
# NVRTC is supplied by Torch. TileLang also needs NVIDIA's matching
# compiler and CCCL headers, without a system CUDA Toolkit.
requirements += [
"nvidia-cuda-nvcc-cu12==12.8.93",
"nvidia-cuda-cccl-cu12==12.8.90",
]
modified_requirements = []
for req in requirements:
if pin_torch_cu128:
req = torch_cu128_urls.get(req.split("#", 1)[0].strip(), req)
if "vllm-flash-attn" in req and cuda_major != "12":
# vllm-flash-attn is built only for CUDA 12.x.
# Skip for other versions.
continue
if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
# [cu13] extra is the default; strip it on CUDA 12 builds.
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._moe_C"))
ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
# Optional since this doesn't get built (produce an .so file). This is just
# copying the relevant .py files from the source repository.
ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
if _is_cuda():
if _cuda_arch_contains(7, 0) or _cuda_arch_contains(7, 2):
ext_modules.append(CMakeExtension(name="vllm._sm70_sampler_C"))
# These extensions use pybind11/libtorch_python and therefore require
# the interpreter-specific CPython ABI suffix emitted by CMake.
ext_modules.append(
CMakeExtension(name="vllm._sm70_qsa_device_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._sm70_qsa_indexer_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._sm70_router_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._sm70_exact_reduce_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._gguf_expert_cache_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._h3_w8a16_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._h3_flashinfer_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._h3_flashattn_C", py_limited_api=False)
)
ext_modules.append(
CMakeExtension(name="vllm._sm70_sparse_attention_C", py_limited_api=False)
)
build_sm70_fa2 = (
_cuda_arch_contains(7, 0) or _cuda_arch_contains(7, 2)
) and not _cuda_arch_at_least(8, 0)
if _cuda_arch_at_least(8, 0) or build_sm70_fa2:
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
if _cuda_arch_contains(7, 5):
# Turing FA2 from a pinned fork, next to the regular library
# (cmake/external_projects/vllm_flash_attn_sm75.cmake).
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C_sm75"))
if _cuda_arch_at_least(8, 0):
if _cuda_arch_at_least(9, 0) and (
USE_PRECOMPILED_EXTENSIONS
or (CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3"))
):
# FA3 requires CUDA 12.3 or later
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
# FA4 CuteDSL - Python-only component for FA4's cute DSL support
# Optional since this doesn't produce a .so file, just copies Python files
ext_modules.append(
CMakeExtension(
name="vllm.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True
)
)
elif not build_sm70_fa2:
logger.info(
"Skipping vllm-flash-attn extension targets for TORCH_CUDA_ARCH_LIST=%s",
os.environ.get("TORCH_CUDA_ARCH_LIST"),
)
if USE_PRECOMPILED_EXTENSIONS or (
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
):
# FlashMLA requires CUDA 12.9 or later
# Optional since this doesn't get built (produce an .so file) when
# not targeting a hopper system
if _cuda_arch_at_least(9, 0):
ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
ext_modules.append(
CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
)
else:
logger.info(
"Skipping FlashMLA extension targets for TORCH_CUDA_ARCH_LIST=%s",
os.environ.get("TORCH_CUDA_ARCH_LIST"),
)
if (
envs.VLLM_USE_PRECOMPILED
or (CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3"))
) and _cuda_arch_at_least(9, 0):
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", 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():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm._C_gguf"))
package_data = {
"flash_qla": [
"LICENSE",
"ops/gated_delta_rule/chunk/sm70/csrc/*.cu",
"ops/gated_delta_rule/chunk/sm70/*.so",
],
"flash_attn_v100": [
"*.so",
],
"vllm": [
"py.typed",
"assets/*.pt",
"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",
# DeepGEMM JIT include headers (vendored via cmake)
"third_party/deep_gemm/include/**/*.cuh",
"third_party/deep_gemm/include/**/*.h",
"third_party/deep_gemm/include/**/*.hpp",
"sm70_profiles/*.json",
],
}
# If using precompiled artifacts, extract and patch package_data in advance.
if USE_PRECOMPILED_RUST_FRONTEND:
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 the rust frontend binary is already present in the source tree (e.g.,
# pre-built in a separate Docker build stage), ship it as-is.
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
vllm_files = package_data.setdefault("vllm", [])
if "vllm-rs" not in vllm_files:
vllm_files.append("vllm-rs")
# A focused CMake source build may prepare the new GGUF extension while the
# unchanged base operators come from a precompiled wheel. Ship that normally
# installed module in the final wheel, rather than a task-cache sidecar.
if USE_PRECOMPILED_EXTENSIONS and (ROOT_DIR / "vllm/_C_gguf.abi3.so").is_file():
package_data.setdefault("vllm", []).append("_C_gguf.abi3.so")
package_data.setdefault("vllm", []).append("third_party/gguf_native/*")
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():
cmdclass["build_rust"] = precompiled_build_rust
cmdclass["bdist_wheel"] = vllm_bdist_wheel
# Rust frontend binary, built via setuptools-rust and installed into the
# package directory alongside the Python modules.
# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
# requires extra work on using precompiled binaries.
rust_extensions = [
RustExtension(
target="vllm.vllm-rs",
path="rust/src/cmd/Cargo.toml",
args=["--bin", "vllm-rs"],
features=["native-tls-vendored"],
binding=Binding.Exec,
optional=not should_require_rust_frontend(),
),
]
setup(
# static metadata should rather go in pyproject.toml
version=get_vllm_version(),
packages=(find_packages(include=["vllm*", "flash_qla*"]) + ["flash_attn_v100"]),
scripts=[
"scripts/serve_qwen38_27b_nvfp4_v100.sh",
"scripts/serve_flash_next_nvfp4_v100.sh",
],
package_dir={
"flash_attn_v100": "flash-attention-v100/flash_attn_v100",
},
ext_modules=ext_modules,
rust_extensions=rust_extensions,
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": [
"zentorch-weekly==5.2.1.dev20260408"
], # Zentorch has weekly releases. This pulls the known-good version.
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.12.1"],
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
"instanttensor": ["instanttensor >= 0.1.5"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.16.1"],
"audio": [
"av",
"scipy",
"soundfile",
"mistral_common[audio]",
], # Required for audio processing
"image": ["diffusers==0.40.0", "accelerate>=1.12", "nvidia-ml-py"],
"video": [
"diffusers==0.40.0",
"av>=14",
"imageio>=2.37.2",
"imageio-ffmpeg>=0.6",
"soundfile>=0.13",
"scipy",
"einops",
"omegaconf",
"accelerate>=1.12",
"nvidia-ml-py",
],
"flashinfer": [], # Kept for backwards compatibility
# Optional deps for Helion kernel development
# NOTE: When updating helion version, also update CI files:
# - .buildkite/test_areas/kernels.yaml
# - .buildkite/test-amd.yaml
"helion": ["helion==1.4.0"],
# Optional deps for gRPC server (vllm serve --grpc)
"grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"],
# Optional deps for OpenTelemetry tracing
"otel": [
"opentelemetry-sdk>=1.26.0",
"opentelemetry-api>=1.26.0",
"opentelemetry-exporter-otlp>=1.26.0",
"opentelemetry-semantic-conventions-ai>=0.4.1",
],
},
cmdclass=cmdclass,
package_data=package_data,
)
Reference in new issue View Git Blame Copy Permalink