| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import ctypes |
| 5 | import importlib.util |
| 6 | import json |
| 7 | import logging |
| 8 | import os |
| 9 | import re |
| 10 | import shutil |
| 11 | import subprocess |
| 12 | import sys |
| 13 | import sysconfig |
| 14 | from pathlib import Path |
| 15 | from shutil import which |
| 16 | |
| 17 | import torch |
| 18 | from packaging.version import Version, parse |
| 19 | from setuptools import Extension, setup |
| 20 | from setuptools.command.build_ext import build_ext |
| 21 | from setuptools_rust.build import build_rust |
| 22 | from setuptools_scm import get_version |
| 23 | from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME |
| 24 | |
| 25 | # Select vLLM release tags, excluding crate tags such as "proto-v0.3.0". |
| 26 | VLLM_GIT_DESCRIBE_COMMAND = [ |
| 27 | "git", |
| 28 | "describe", |
| 29 | "--dirty", |
| 30 | "--tags", |
| 31 | "--long", |
| 32 | "--abbrev=40", |
| 33 | "--match", |
| 34 | "v[0-9]*", |
| 35 | ] |
| 36 | |
| 37 | |
| 38 | def load_module_from_path(module_name, path): |
| 39 | spec = importlib.util.spec_from_file_location(module_name, path) |
| 40 | module = importlib.util.module_from_spec(spec) |
| 41 | sys.modules[module_name] = module |
| 42 | spec.loader.exec_module(module) |
| 43 | return module |
| 44 | |
| 45 | |
| 46 | ROOT_DIR = Path(__file__).parent |
| 47 | logger = logging.getLogger(__name__) |
| 48 | |
| 49 | PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs" |
| 50 | # setuptools-rust installs PyO3 artifacts as `<module>.<ext-suffix>`, where the |
| 51 | # suffix ends with `.so` on Linux and macOS alike (e.g. `_rust_foo.abi3.so`). |
| 52 | PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX = re.compile(r"vllm/_rust_[^/]*\.so$") |
| 53 | |
| 54 | # cannot import envs directly because it depends on vllm, |
| 55 | # which is not installed yet |
| 56 | envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py")) |
| 57 | rust_build = load_module_from_path( |
| 58 | "rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py") |
| 59 | ) |
| 60 | |
| 61 | VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE |
| 62 | USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED |
| 63 | # VLLM_USE_PRECOMPILED implies precompiled rust frontend too. |
| 64 | USE_PRECOMPILED_RUST_FRONTEND = ( |
| 65 | envs.VLLM_USE_PRECOMPILED or envs.VLLM_USE_PRECOMPILED_RUST |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def should_require_rust_frontend() -> bool: |
| 70 | value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "") |
| 71 | return value.lower() not in ("", "0", "false", "no") |
| 72 | |
| 73 | |
| 74 | def get_precompiled_rust_extension_paths() -> list[Path]: |
| 75 | return sorted((ROOT_DIR / "vllm").glob("_rust_*.so")) |
| 76 | |
| 77 | |
| 78 | def get_missing_precompiled_rust_extension_modules() -> list[str]: |
| 79 | present = { |
| 80 | path.name.split(".", 1)[0] for path in get_precompiled_rust_extension_paths() |
| 81 | } |
| 82 | return [ |
| 83 | module_name |
| 84 | for module_name in rust_build.rust_py_extension_module_names() |
| 85 | if module_name not in present |
| 86 | ] |
| 87 | |
| 88 | |
| 89 | def has_precompiled_rust_extensions() -> bool: |
| 90 | return not get_missing_precompiled_rust_extension_modules() |
| 91 | |
| 92 | |
| 93 | def is_metadata_only_build() -> bool: |
| 94 | return bool({"egg_info", "dist_info"}.intersection(sys.argv[1:])) |
| 95 | |
| 96 | |
| 97 | if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu": |
| 98 | logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS") |
| 99 | VLLM_TARGET_DEVICE = "cpu" |
| 100 | elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")): |
| 101 | logger.warning( |
| 102 | "vLLM only supports Linux platform (including WSL) and MacOS." |
| 103 | "Building on %s, " |
| 104 | "so vLLM may not be able to run correctly", |
| 105 | sys.platform, |
| 106 | ) |
| 107 | VLLM_TARGET_DEVICE = "empty" |
| 108 | elif sys.platform.startswith("linux") and os.getenv("VLLM_TARGET_DEVICE") is None: |
| 109 | if torch.version.hip is not None: |
| 110 | VLLM_TARGET_DEVICE = "rocm" |
| 111 | logger.info("Auto-detected ROCm") |
| 112 | elif torch.version.xpu is not None: |
| 113 | VLLM_TARGET_DEVICE = "xpu" |
| 114 | logger.info("Auto-detected XPU") |
| 115 | elif torch.version.cuda is not None: |
| 116 | VLLM_TARGET_DEVICE = "cuda" |
| 117 | logger.info("Auto-detected CUDA") |
| 118 | else: |
| 119 | VLLM_TARGET_DEVICE = "cpu" |
| 120 | |
| 121 | |
| 122 | def is_sccache_available() -> bool: |
| 123 | return which("sccache") is not None and not bool( |
| 124 | int(os.getenv("VLLM_DISABLE_SCCACHE", "0")) |
| 125 | ) |
| 126 | |
| 127 | |
| 128 | def is_ccache_available() -> bool: |
| 129 | return which("ccache") is not None |
| 130 | |
| 131 | |
| 132 | def is_ninja_available() -> bool: |
| 133 | return which("ninja") is not None |
| 134 | |
| 135 | |
| 136 | def is_freethreaded(): |
| 137 | return bool(sysconfig.get_config_var("Py_GIL_DISABLED")) |
| 138 | |
| 139 | |
| 140 | def should_bundle_tcmalloc() -> bool: |
| 141 | import platform |
| 142 | |
| 143 | return ( |
| 144 | VLLM_TARGET_DEVICE == "cpu" |
| 145 | and sys.platform.startswith("linux") |
| 146 | and platform.machine() in ("aarch64", "x86_64", "s390x") |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | def find_tcmalloc() -> Path | None: |
| 151 | try: |
| 152 | # get all shared libs the dynamic loader knows about |
| 153 | output = subprocess.check_output( |
| 154 | ["ldconfig", "-p"], |
| 155 | text=True, |
| 156 | stderr=subprocess.DEVNULL, |
| 157 | ) |
| 158 | except Exception: |
| 159 | return None |
| 160 | |
| 161 | # search for libtcmalloc and libtcmalloc_minimal |
| 162 | for library_pattern in ( |
| 163 | r"\blibtcmalloc_minimal\.so\.(\d+)\b", |
| 164 | r"\blibtcmalloc\.so\.(\d+)\b", |
| 165 | ): |
| 166 | candidates: list[tuple[int, Path]] = [] |
| 167 | for line in output.splitlines(): |
| 168 | match = re.search(library_pattern, line) |
| 169 | if match is None or "=>" not in line: |
| 170 | continue |
| 171 | candidate = Path(line.split("=>")[1].strip()) |
| 172 | if candidate.exists(): |
| 173 | candidates.append((int(match.group(1)), candidate)) |
| 174 | |
| 175 | if candidates: |
| 176 | # if multiple candidates are found, pick the one with the highest |
| 177 | # version number |
| 178 | return max(candidates, key=lambda item: item[0])[1] |
| 179 | |
| 180 | return None |
| 181 | |
| 182 | |
| 183 | def bundle_tcmalloc(build_lib: str) -> None: |
| 184 | tcmalloc_library = find_tcmalloc() |
| 185 | if tcmalloc_library is None: |
| 186 | logger.warning( |
| 187 | "Failed to locate tcmalloc. For best performance, " |
| 188 | "please install tcmalloc (e.g. `sudo apt-get " |
| 189 | "install -y --no-install-recommends libtcmalloc-minimal4`)" |
| 190 | ) |
| 191 | return |
| 192 | |
| 193 | bundle_dir = os.path.join(build_lib, "vllm", "libs") |
| 194 | os.makedirs(bundle_dir, exist_ok=True) |
| 195 | bundle_path = os.path.join(bundle_dir, tcmalloc_library.name) |
| 196 | shutil.copy2(tcmalloc_library, bundle_path) |
| 197 | logger.info("Bundled tcmalloc into wheel: %s", bundle_path) |
| 198 | |
| 199 | |
| 200 | class CMakeExtension(Extension): |
| 201 | def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None: |
| 202 | super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa) |
| 203 | self.cmake_lists_dir = os.path.abspath(cmake_lists_dir) |
| 204 | |
| 205 | |
| 206 | class cmake_build_ext(build_ext): |
| 207 | # A dict of extension directories that have been configured. |
| 208 | did_config: dict[str, bool] = {} |
| 209 | |
| 210 | # |
| 211 | # Determine number of compilation jobs and optionally nvcc compile threads. |
| 212 | # |
| 213 | def compute_num_jobs(self): |
| 214 | # `num_jobs` is either the value of the MAX_JOBS environment variable |
| 215 | # (if defined) or the number of CPUs available. |
| 216 | num_jobs = envs.MAX_JOBS |
| 217 | if num_jobs is not None: |
| 218 | num_jobs = int(num_jobs) |
| 219 | logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs) |
| 220 | else: |
| 221 | try: |
| 222 | # os.sched_getaffinity() isn't universally available, so fall |
| 223 | # back to os.cpu_count() if we get an error here. |
| 224 | num_jobs = len(os.sched_getaffinity(0)) |
| 225 | except AttributeError: |
| 226 | num_jobs = os.cpu_count() |
| 227 | |
| 228 | nvcc_threads = None |
| 229 | if _is_cuda() and CUDA_HOME is not None: |
| 230 | try: |
| 231 | nvcc_version = get_nvcc_cuda_version() |
| 232 | if nvcc_version >= Version("11.2"): |
| 233 | # `nvcc_threads` is either the value of the NVCC_THREADS |
| 234 | # environment variable (if defined) or 1. |
| 235 | # when it is set, we reduce `num_jobs` to avoid |
| 236 | # overloading the system. |
| 237 | nvcc_threads = envs.NVCC_THREADS |
| 238 | if nvcc_threads is not None: |
| 239 | nvcc_threads = int(nvcc_threads) |
| 240 | logger.info( |
| 241 | "Using NVCC_THREADS=%d as the number of nvcc threads.", |
| 242 | nvcc_threads, |
| 243 | ) |
| 244 | else: |
| 245 | nvcc_threads = 1 |
| 246 | num_jobs = max(1, num_jobs // nvcc_threads) |
| 247 | except Exception as e: |
| 248 | logger.warning("Failed to get NVCC version: %s", e) |
| 249 | |
| 250 | return num_jobs, nvcc_threads |
| 251 | |
| 252 | # |
| 253 | # Perform cmake configuration for a single extension. |
| 254 | # |
| 255 | def configure(self, ext: CMakeExtension) -> None: |
| 256 | # If we've already configured using the CMakeLists.txt for |
| 257 | # this extension, exit early. |
| 258 | if ext.cmake_lists_dir in cmake_build_ext.did_config: |
| 259 | return |
| 260 | |
| 261 | cmake_build_ext.did_config[ext.cmake_lists_dir] = True |
| 262 | |
| 263 | # Select the build type. |
| 264 | # Note: optimization level + debug info are set by the build type |
| 265 | default_cfg = "Debug" if self.debug else "RelWithDebInfo" |
| 266 | cfg = envs.CMAKE_BUILD_TYPE or default_cfg |
| 267 | |
| 268 | cmake_args = [ |
| 269 | "-DCMAKE_BUILD_TYPE={}".format(cfg), |
| 270 | "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE), |
| 271 | ] |
| 272 | |
| 273 | verbose = envs.VERBOSE |
| 274 | if verbose: |
| 275 | cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"] |
| 276 | |
| 277 | if is_sccache_available(): |
| 278 | cmake_args += [ |
| 279 | "-DCMAKE_C_COMPILER_LAUNCHER=sccache", |
| 280 | "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache", |
| 281 | "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache", |
| 282 | "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache", |
| 283 | ] |
| 284 | elif is_ccache_available(): |
| 285 | cmake_args += [ |
| 286 | "-DCMAKE_C_COMPILER_LAUNCHER=ccache", |
| 287 | "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache", |
| 288 | "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache", |
| 289 | "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache", |
| 290 | ] |
| 291 | |
| 292 | # Pass the python executable to cmake so it can find an exact |
| 293 | # match. |
| 294 | cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)] |
| 295 | |
| 296 | # Pass the python path to cmake so it can reuse the build dependencies |
| 297 | # on subsequent calls to python. |
| 298 | cmake_args += ["-DVLLM_PYTHON_PATH={}".format(":".join(sys.path))] |
| 299 | |
| 300 | # Override the base directory for FetchContent downloads to $ROOT/.deps |
| 301 | # This allows sharing dependencies between profiles, |
| 302 | # and plays more nicely with sccache. |
| 303 | # To override this, set the FETCHCONTENT_BASE_DIR environment variable. |
| 304 | fc_base_dir = os.path.join(ROOT_DIR, ".deps") |
| 305 | fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir) |
| 306 | cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)] |
| 307 | |
| 308 | # |
| 309 | # Setup parallelism and build tool |
| 310 | # |
| 311 | num_jobs, nvcc_threads = self.compute_num_jobs() |
| 312 | |
| 313 | if nvcc_threads: |
| 314 | cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)] |
| 315 | |
| 316 | if is_ninja_available(): |
| 317 | build_tool = ["-G", "Ninja"] |
| 318 | cmake_args += [ |
| 319 | "-DCMAKE_JOB_POOL_COMPILE:STRING=compile", |
| 320 | "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs), |
| 321 | ] |
| 322 | else: |
| 323 | # Default build tool to whatever cmake picks. |
| 324 | build_tool = [] |
| 325 | # Make sure we use the nvcc from CUDA_HOME |
| 326 | if _is_cuda() and CUDA_HOME is not None: |
| 327 | cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"] |
| 328 | elif _is_hip() and ROCM_HOME is not None: |
| 329 | cmake_args += [f"-DROCM_PATH={ROCM_HOME}"] |
| 330 | |
| 331 | other_cmake_args = os.environ.get("CMAKE_ARGS") |
| 332 | if other_cmake_args: |
| 333 | cmake_args += other_cmake_args.split() |
| 334 | |
| 335 | subprocess.check_call( |
| 336 | ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args], |
| 337 | cwd=self.build_temp, |
| 338 | ) |
| 339 | |
| 340 | def build_extensions(self) -> None: |
| 341 | # Ensure that CMake is present and working |
| 342 | try: |
| 343 | subprocess.check_output(["cmake", "--version"]) |
| 344 | except OSError as e: |
| 345 | raise RuntimeError("Cannot find CMake executable") from e |
| 346 | |
| 347 | # Create build directory if it does not exist. |
| 348 | if not os.path.exists(self.build_temp): |
| 349 | os.makedirs(self.build_temp) |
| 350 | |
| 351 | targets = [] |
| 352 | |
| 353 | def target_name(s: str) -> str: |
| 354 | return s.removeprefix("vllm.").removeprefix("vllm_flash_attn.") |
| 355 | |
| 356 | # Build all the extensions |
| 357 | for ext in self.extensions: |
| 358 | self.configure(ext) |
| 359 | targets.append(target_name(ext.name)) |
| 360 | |
| 361 | num_jobs, _ = self.compute_num_jobs() |
| 362 | |
| 363 | build_args = [ |
| 364 | "--build", |
| 365 | ".", |
| 366 | f"-j={num_jobs}", |
| 367 | *[f"--target={name}" for name in targets], |
| 368 | ] |
| 369 | |
| 370 | subprocess.check_call(["cmake", *build_args], cwd=self.build_temp) |
| 371 | |
| 372 | # Install the libraries |
| 373 | for ext in self.extensions: |
| 374 | # Install the extension into the proper location |
| 375 | outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute() |
| 376 | |
| 377 | # Skip if the install directory is the same as the build directory |
| 378 | if outdir == self.build_temp: |
| 379 | continue |
| 380 | |
| 381 | # CMake appends the extension prefix to the install path, |
| 382 | # and outdir already contains that prefix, so we need to remove it. |
| 383 | prefix = outdir |
| 384 | for _ in range(ext.name.count(".")): |
| 385 | prefix = prefix.parent |
| 386 | |
| 387 | # prefix here should actually be the same for all components |
| 388 | install_args = [ |
| 389 | "cmake", |
| 390 | "--install", |
| 391 | ".", |
| 392 | "--prefix", |
| 393 | prefix, |
| 394 | "--component", |
| 395 | target_name(ext.name), |
| 396 | ] |
| 397 | subprocess.check_call(install_args, cwd=self.build_temp) |
| 398 | |
| 399 | def run(self): |
| 400 | # First, run the standard build_ext command to compile the extensions |
| 401 | super().run() |
| 402 | |
| 403 | # bundle tcmalloc into CPU wheels for best OOB perf |
| 404 | if should_bundle_tcmalloc(): |
| 405 | bundle_tcmalloc(self.build_lib) |
| 406 | |
| 407 | # copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current |
| 408 | # directory so that they can be included in the editable build |
| 409 | import glob |
| 410 | |
| 411 | files = glob.glob( |
| 412 | os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "**", "*.py"), |
| 413 | recursive=True, |
| 414 | ) |
| 415 | for file in files: |
| 416 | dst_file = os.path.join( |
| 417 | "vllm/vllm_flash_attn", file.split("vllm/vllm_flash_attn/")[-1] |
| 418 | ) |
| 419 | print(f"Copying {file} to {dst_file}") |
| 420 | os.makedirs(os.path.dirname(dst_file), exist_ok=True) |
| 421 | self.copy_file(file, dst_file) |
| 422 | |
| 423 | if _is_cuda() or _is_hip(): |
| 424 | # copy vllm/third_party/triton_kernels/**/*.py from self.build_lib |
| 425 | # to current directory so that they can be included in the editable |
| 426 | # build |
| 427 | print( |
| 428 | f"Copying {self.build_lib}/vllm/third_party/triton_kernels " |
| 429 | "to vllm/third_party/triton_kernels" |
| 430 | ) |
| 431 | shutil.copytree( |
| 432 | f"{self.build_lib}/vllm/third_party/triton_kernels", |
| 433 | "vllm/third_party/triton_kernels", |
| 434 | dirs_exist_ok=True, |
| 435 | ) |
| 436 | |
| 437 | if _is_cuda(): |
| 438 | # copy vendored deep_gemm package from build_lib to source tree |
| 439 | # for editable installs |
| 440 | deep_gemm_build = os.path.join( |
| 441 | self.build_lib, "vllm", "third_party", "deep_gemm" |
| 442 | ) |
| 443 | if os.path.exists(deep_gemm_build): |
| 444 | print(f"Copying {deep_gemm_build} to vllm/third_party/deep_gemm") |
| 445 | shutil.copytree( |
| 446 | deep_gemm_build, |
| 447 | "vllm/third_party/deep_gemm", |
| 448 | dirs_exist_ok=True, |
| 449 | ) |
| 450 | |
| 451 | # copy vendored fmha_sm100 package from build_lib to source tree |
| 452 | # for editable installs |
| 453 | fmha_sm100_build = os.path.join( |
| 454 | self.build_lib, "vllm", "third_party", "fmha_sm100" |
| 455 | ) |
| 456 | if os.path.exists(fmha_sm100_build): |
| 457 | print(f"Copying {fmha_sm100_build} to vllm/third_party/fmha_sm100") |
| 458 | shutil.copytree( |
| 459 | fmha_sm100_build, |
| 460 | "vllm/third_party/fmha_sm100", |
| 461 | dirs_exist_ok=True, |
| 462 | ) |
| 463 | |
| 464 | tml_fa4_build = os.path.join( |
| 465 | self.build_lib, "vllm", "third_party", "tml_fa4" |
| 466 | ) |
| 467 | if os.path.exists(tml_fa4_build): |
| 468 | print(f"Copying {tml_fa4_build} to vllm/third_party/tml_fa4") |
| 469 | shutil.copytree( |
| 470 | tml_fa4_build, |
| 471 | "vllm/third_party/tml_fa4", |
| 472 | dirs_exist_ok=True, |
| 473 | ) |
| 474 | |
| 475 | |
| 476 | class precompiled_build_ext(build_ext): |
| 477 | """Disables extension building when using precompiled binaries.""" |
| 478 | |
| 479 | def run(self) -> None: |
| 480 | return |
| 481 | |
| 482 | def build_extensions(self) -> None: |
| 483 | print("Skipping build_ext: using precompiled extensions.") |
| 484 | return |
| 485 | |
| 486 | |
| 487 | class precompiled_build_rust(build_rust): |
| 488 | """Skips local Rust builds when all precompiled Rust artifacts are present.""" |
| 489 | |
| 490 | def run(self) -> None: |
| 491 | missing = [] |
| 492 | if not PRECOMPILED_RUST_FRONTEND_PATH.exists(): |
| 493 | missing.append(str(PRECOMPILED_RUST_FRONTEND_PATH)) |
| 494 | missing_rust_extensions = get_missing_precompiled_rust_extension_modules() |
| 495 | if missing_rust_extensions: |
| 496 | missing.extend( |
| 497 | str(ROOT_DIR / "vllm" / f"{module_name}*.so") |
| 498 | for module_name in missing_rust_extensions |
| 499 | ) |
| 500 | |
| 501 | if not missing: |
| 502 | logger.info( |
| 503 | "Skipping local Rust build: using precompiled %s and %s", |
| 504 | PRECOMPILED_RUST_FRONTEND_PATH, |
| 505 | get_precompiled_rust_extension_paths(), |
| 506 | ) |
| 507 | return |
| 508 | |
| 509 | logger.warning( |
| 510 | "Precompiled wheel did not provide all Rust artifacts (%s); " |
| 511 | "falling back to local Rust build.", |
| 512 | ", ".join(missing), |
| 513 | ) |
| 514 | super().run() |
| 515 | |
| 516 | |
| 517 | class precompiled_wheel_utils: |
| 518 | """Extracts libraries and other files from an existing wheel.""" |
| 519 | |
| 520 | @staticmethod |
| 521 | def fetch_metadata_for_variant( |
| 522 | commit: str, |
| 523 | variant: str | None, |
| 524 | *, |
| 525 | rocm: bool = False, |
| 526 | ) -> tuple[list[dict], str]: |
| 527 | """Fetches metadata for a specific variant of the precompiled wheel. |
| 528 | |
| 529 | For non-ROCm, fetches vllm metadata. |
| 530 | |
| 531 | For ROCm, discovers all first-level packages and combines their |
| 532 | metadata into a single list. |
| 533 | """ |
| 534 | import json |
| 535 | from html.parser import HTMLParser |
| 536 | from urllib.request import urlopen |
| 537 | |
| 538 | variant_dir = f"{variant}/" if variant is not None else "" |
| 539 | |
| 540 | if not rocm: |
| 541 | # Keep original behavior |
| 542 | repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/" |
| 543 | meta_url = repo_url + "metadata.json" |
| 544 | print(f"Trying to fetch nightly build metadata from {meta_url}") |
| 545 | with urlopen(meta_url) as resp: |
| 546 | wheels = json.loads(resp.read().decode("utf-8")) |
| 547 | |
| 548 | return wheels, repo_url |
| 549 | |
| 550 | # ROCm: discover all packages under the variant directory. |
| 551 | repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}" |
| 552 | |
| 553 | class LinkParser(HTMLParser): |
| 554 | def __init__(self): |
| 555 | super().__init__() |
| 556 | self.links: list[str] = [] |
| 557 | |
| 558 | def handle_starttag(self, tag, attrs): |
| 559 | if tag == "a": |
| 560 | href = dict(attrs).get("href") |
| 561 | if href: |
| 562 | self.links.append(href) |
| 563 | |
| 564 | with urlopen(repo_url) as resp: |
| 565 | parser = LinkParser() |
| 566 | parser.feed(resp.read().decode("utf-8")) |
| 567 | |
| 568 | packages = [ |
| 569 | href.rstrip("/") |
| 570 | for href in parser.links |
| 571 | if href.endswith("/") |
| 572 | and href not in ("../", "./", "/") |
| 573 | and "/" not in href.rstrip("/") |
| 574 | ] |
| 575 | |
| 576 | wheels: list[dict] = [] |
| 577 | |
| 578 | for package in packages: |
| 579 | meta_url = f"{repo_url}{package}/metadata.json" |
| 580 | print(f"Trying to fetch nightly build metadata from {meta_url}") |
| 581 | with urlopen(meta_url) as resp: |
| 582 | package_wheels = json.loads(resp.read().decode("utf-8")) |
| 583 | wheels.extend(package_wheels) |
| 584 | return wheels, repo_url |
| 585 | |
| 586 | @staticmethod |
| 587 | def is_rocm_system() -> bool: |
| 588 | """Detect ROCm without relying on torch (for build environment).""" |
| 589 | if os.getenv("ROCM_PATH"): |
| 590 | return True |
| 591 | if os.path.isdir("/opt/rocm"): |
| 592 | return True |
| 593 | if which("rocminfo") is not None: |
| 594 | return True |
| 595 | try: |
| 596 | import torch |
| 597 | |
| 598 | return torch.version.hip is not None |
| 599 | except ImportError: |
| 600 | return False |
| 601 | |
| 602 | @staticmethod |
| 603 | def detect_system_cuda_variant() -> str: |
| 604 | """Auto-detect CUDA variant from torch, nvidia-smi, or env default.""" |
| 605 | # Map CUDA major version to hosted wheel variants on wheels.vllm.ai |
| 606 | supported = {12: "cu129", 13: "cu130"} |
| 607 | |
| 608 | # Respect explicitly set VLLM_MAIN_CUDA_VERSION |
| 609 | if envs.is_set("VLLM_MAIN_CUDA_VERSION"): |
| 610 | v = envs.VLLM_MAIN_CUDA_VERSION |
| 611 | print(f"Using VLLM_MAIN_CUDA_VERSION={v}") |
| 612 | return "cu" + v.replace(".", "")[:3] |
| 613 | |
| 614 | # Try torch.version.cuda |
| 615 | cuda_version = None |
| 616 | try: |
| 617 | import torch |
| 618 | |
| 619 | cuda_version = torch.version.cuda |
| 620 | except Exception: |
| 621 | pass |
| 622 | |
| 623 | # Try nvidia-smi |
| 624 | if not cuda_version: |
| 625 | try: |
| 626 | out = subprocess.run( |
| 627 | ["nvidia-smi"], capture_output=True, text=True, timeout=10 |
| 628 | ) |
| 629 | if m := re.search(r"CUDA Version:\s*(\d+\.\d+)", out.stdout): |
| 630 | cuda_version = m.group(1) |
| 631 | except Exception: |
| 632 | pass |
| 633 | |
| 634 | # Fall back to default |
| 635 | if not cuda_version: |
| 636 | cuda_version = envs.VLLM_MAIN_CUDA_VERSION |
| 637 | |
| 638 | # Map to supported variant |
| 639 | major = int(cuda_version.split(".")[0]) |
| 640 | variant = supported.get(major, supported[max(supported)]) |
| 641 | print(f"Detected CUDA {cuda_version}, using variant {variant}") |
| 642 | return variant |
| 643 | |
| 644 | @staticmethod |
| 645 | def rocm_version_to_variant(rocm_version: str) -> str: |
| 646 | """Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723.""" |
| 647 | return "rocm" + rocm_version.replace(".", "") |
| 648 | |
| 649 | @staticmethod |
| 650 | def detect_system_rocm_variant() -> str | None: |
| 651 | """Auto-detect the ROCm wheel variant from the installed ROCm stack.""" |
| 652 | rocm_version = get_rocm_version() |
| 653 | if not rocm_version: |
| 654 | try: |
| 655 | import torch |
| 656 | |
| 657 | rocm_version = torch.version.hip |
| 658 | except Exception: |
| 659 | pass |
| 660 | if not rocm_version: |
| 661 | return None |
| 662 | variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version) |
| 663 | print(f"Detected ROCm {rocm_version}, using variant {variant}") |
| 664 | return variant |
| 665 | |
| 666 | @staticmethod |
| 667 | def fetch_available_rocm_variants(commit: str) -> list[str]: |
| 668 | """List ROCm wheel variants published for a commit on wheels.vllm.ai.""" |
| 669 | from urllib.request import urlopen |
| 670 | |
| 671 | index_url = f"https://wheels.vllm.ai/rocm/{commit}/" |
| 672 | print(f"Fetching available ROCm variants from {index_url}") |
| 673 | try: |
| 674 | with urlopen(index_url) as resp: |
| 675 | html = resp.read().decode("utf-8") |
| 676 | except Exception as e: |
| 677 | logger.warning( |
| 678 | "Failed to fetch ROCm variant index for commit %s: %s", commit, e |
| 679 | ) |
| 680 | return [] |
| 681 | variants = sorted(set(re.findall(r"rocm\d+", html))) |
| 682 | print(f"Available ROCm variants for commit {commit}: {variants}") |
| 683 | return variants |
| 684 | |
| 685 | @staticmethod |
| 686 | def resolve_rocm_wheel_variant( |
| 687 | commit: str, variant_override: str | None |
| 688 | ) -> str | None: |
| 689 | """Resolve a ROCm wheel variant for a commit from wheels.vllm.ai.""" |
| 690 | env_variant = precompiled_wheel_utils.detect_system_rocm_variant() |
| 691 | available = precompiled_wheel_utils.fetch_available_rocm_variants(commit) |
| 692 | |
| 693 | if variant_override is not None: |
| 694 | if env_variant and variant_override != env_variant: |
| 695 | logger.warning( |
| 696 | "VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the " |
| 697 | "detected environment ROCm variant %s; refusing to use a " |
| 698 | "different ROCm patch wheel", |
| 699 | variant_override, |
| 700 | env_variant, |
| 701 | ) |
| 702 | return None |
| 703 | if available and variant_override not in available: |
| 704 | logger.warning( |
| 705 | "Requested ROCm variant %s is not available for commit %s " |
| 706 | "(available: %s)", |
| 707 | variant_override, |
| 708 | commit, |
| 709 | available, |
| 710 | ) |
| 711 | return None |
| 712 | return variant_override |
| 713 | |
| 714 | if env_variant is None: |
| 715 | if available: |
| 716 | print( |
| 717 | "Could not detect ROCm version from the environment; " |
| 718 | f"available variants for commit {commit}: {available}" |
| 719 | ) |
| 720 | return None |
| 721 | |
| 722 | if env_variant in available: |
| 723 | return env_variant |
| 724 | |
| 725 | logger.warning( |
| 726 | "Environment ROCm variant %s is not available for commit %s " |
| 727 | "(available: %s). Precompiled wheels may not be compatible.", |
| 728 | env_variant, |
| 729 | commit, |
| 730 | available, |
| 731 | ) |
| 732 | return None |
| 733 | |
| 734 | @staticmethod |
| 735 | def warn_if_rocm_torch_version_mismatch( |
| 736 | wheels: list[dict], repo_url: str, arch: str |
| 737 | ) -> None: |
| 738 | """Warn if installed torch differs from the custom ROCm build on |
| 739 | wheels.vllm.ai and suggest the correct install command.""" |
| 740 | try: |
| 741 | installed = torch.__version__ |
| 742 | except Exception: |
| 743 | return |
| 744 | |
| 745 | def _wheel_version(pkg: str) -> str | None: |
| 746 | for w in wheels: |
| 747 | if w.get("package_name") == pkg and arch in w.get("platform_tag", ""): |
| 748 | v = w["version"] |
| 749 | if w.get("variant") and "+" not in v: |
| 750 | v = f"{v}+{w['variant']}" |
| 751 | return v |
| 752 | return None |
| 753 | |
| 754 | expected = _wheel_version("torch") |
| 755 | if expected is None or installed == expected: |
| 756 | return |
| 757 | |
| 758 | pkgs = f"torch=={expected}" |
| 759 | triton_ver = _wheel_version("triton") |
| 760 | if triton_ver: |
| 761 | pkgs += f" triton=={triton_ver}" |
| 762 | |
| 763 | logger.warning( |
| 764 | "Installed PyTorch %s does not match the custom build %s " |
| 765 | "shipped with vLLM ROCm wheels. The ABI may differ from " |
| 766 | "official releases. If you hit extension load errors, " |
| 767 | "reinstall from the vLLM index:\n" |
| 768 | " pip install %s --extra-index-url %s\n" |
| 769 | " uv pip install %s --extra-index-url %s " |
| 770 | "--index-strategy unsafe-best-match", |
| 771 | installed, |
| 772 | expected, |
| 773 | pkgs, |
| 774 | repo_url, |
| 775 | pkgs, |
| 776 | repo_url, |
| 777 | ) |
| 778 | |
| 779 | @staticmethod |
| 780 | def find_local_rocm_wheel() -> str | None: |
| 781 | """Search for a local vllm wheel in common locations.""" |
| 782 | import glob |
| 783 | |
| 784 | for pattern in ["/vllm-workspace/dist/vllm-*.whl", "./dist/vllm-*.whl"]: |
| 785 | wheels = glob.glob(pattern) |
| 786 | if wheels: |
| 787 | return sorted(wheels)[-1] |
| 788 | return None |
| 789 | |
| 790 | @staticmethod |
| 791 | def fetch_wheel_from_pypi_index(index_url: str, package: str = "vllm") -> str: |
| 792 | """Fetch the latest wheel URL from a PyPI-style simple index.""" |
| 793 | import platform |
| 794 | from html.parser import HTMLParser |
| 795 | from urllib.parse import urljoin |
| 796 | from urllib.request import urlopen |
| 797 | |
| 798 | arch = platform.machine() |
| 799 | |
| 800 | class WheelLinkParser(HTMLParser): |
| 801 | def __init__(self): |
| 802 | super().__init__() |
| 803 | self.wheels = [] |
| 804 | |
| 805 | def handle_starttag(self, tag, attrs): |
| 806 | if tag == "a": |
| 807 | for name, value in attrs: |
| 808 | if name == "href" and value.endswith(".whl"): |
| 809 | self.wheels.append(value) |
| 810 | |
| 811 | simple_url = f"{index_url.rstrip('/')}/{package}/" |
| 812 | print(f"Fetching wheel list from {simple_url}") |
| 813 | with urlopen(simple_url) as resp: |
| 814 | html = resp.read().decode("utf-8") |
| 815 | |
| 816 | parser = WheelLinkParser() |
| 817 | parser.feed(html) |
| 818 | |
| 819 | for wheel in reversed(parser.wheels): |
| 820 | if arch in wheel: |
| 821 | if wheel.startswith("http"): |
| 822 | return wheel |
| 823 | return urljoin(simple_url, wheel) |
| 824 | |
| 825 | raise ValueError(f"No compatible wheel found for {arch} at {simple_url}") |
| 826 | |
| 827 | @staticmethod |
| 828 | def determine_wheel_url_rocm() -> tuple[str, str | None]: |
| 829 | """Determine the precompiled wheel for ROCm.""" |
| 830 | # Search for local wheel first |
| 831 | local_wheel = precompiled_wheel_utils.find_local_rocm_wheel() |
| 832 | if local_wheel is not None: |
| 833 | print(f"Found local ROCm wheel: {local_wheel}") |
| 834 | return local_wheel, None |
| 835 | |
| 836 | import platform |
| 837 | |
| 838 | arch = platform.machine() |
| 839 | commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower() |
| 840 | if not commit or len(commit) != 40: |
| 841 | print( |
| 842 | f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}" |
| 843 | ", trying to fetch base commit in main branch" |
| 844 | ) |
| 845 | commit = precompiled_wheel_utils.get_base_commit_in_main_branch() |
| 846 | variant = precompiled_wheel_utils.resolve_rocm_wheel_variant( |
| 847 | commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None) |
| 848 | ) |
| 849 | print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}") |
| 850 | wheels, repo_url = None, None |
| 851 | if variant is not None: |
| 852 | try: |
| 853 | wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant( |
| 854 | commit, variant, rocm=True |
| 855 | ) |
| 856 | precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch( |
| 857 | wheels, repo_url, arch |
| 858 | ) |
| 859 | except Exception as e: |
| 860 | logger.warning( |
| 861 | "Failed to fetch ROCm wheel metadata for variant %s: %s", |
| 862 | variant, |
| 863 | e, |
| 864 | ) |
| 865 | if wheels is not None and repo_url is not None: |
| 866 | from urllib.parse import urljoin |
| 867 | |
| 868 | for wheel in wheels: |
| 869 | if wheel.get("package_name") == "vllm" and arch in wheel.get( |
| 870 | "platform_tag", "" |
| 871 | ): |
| 872 | print(f"Found precompiled wheel metadata: {wheel}") |
| 873 | if "path" not in wheel: |
| 874 | raise ValueError(f"Wheel metadata missing path: {wheel}") |
| 875 | wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"]) |
| 876 | download_filename = wheel.get("filename") |
| 877 | print(f"Using precompiled wheel URL: {wheel_url}") |
| 878 | return wheel_url, download_filename |
| 879 | logger.warning( |
| 880 | "No precompiled vllm wheel found for architecture %s " |
| 881 | "from repo %s. All available wheels: %s", |
| 882 | arch, |
| 883 | repo_url, |
| 884 | wheels, |
| 885 | ) |
| 886 | |
| 887 | # Fall back to AMD's PyPI index |
| 888 | index_url = os.getenv( |
| 889 | "VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple" |
| 890 | ) |
| 891 | print(f"Fetching ROCm precompiled wheel from {index_url}") |
| 892 | wheel_url = precompiled_wheel_utils.fetch_wheel_from_pypi_index(index_url) |
| 893 | download_filename = wheel_url.split("/")[-1].split("#")[0] |
| 894 | print(f"Using ROCm precompiled wheel: {wheel_url}") |
| 895 | return wheel_url, download_filename |
| 896 | |
| 897 | @staticmethod |
| 898 | def determine_wheel_url() -> tuple[str, str | None]: |
| 899 | """Try to determine the precompiled wheel URL or path to use. |
| 900 | The order of preference is: |
| 901 | 1. user-specified wheel location (can be either local or remote, via |
| 902 | VLLM_PRECOMPILED_WHEEL_LOCATION) |
| 903 | 2. user-specified variant (VLLM_PRECOMPILED_WHEEL_VARIANT) from nightly repo |
| 904 | or CUDA variant selected from VLLM_MAIN_CUDA_VERSION, torch, or nvidia-smi |
| 905 | |
| 906 | If downloading from the nightly repo, the commit can be specified via |
| 907 | VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch |
| 908 | is used. |
| 909 | """ |
| 910 | wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None) |
| 911 | if wheel_location is not None: |
| 912 | print(f"Using user-specified precompiled wheel location: {wheel_location}") |
| 913 | return wheel_location, None |
| 914 | else: |
| 915 | # ROCm: resolve wheels from wheels.vllm.ai with environment-matched variant |
| 916 | if precompiled_wheel_utils.is_rocm_system(): |
| 917 | return precompiled_wheel_utils.determine_wheel_url_rocm() |
| 918 | |
| 919 | import platform |
| 920 | |
| 921 | arch = platform.machine() |
| 922 | # try to fetch the wheel metadata from the nightly wheel repo, |
| 923 | # detecting CUDA variant from system if not specified |
| 924 | variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None) |
| 925 | if variant is None: |
| 926 | variant = precompiled_wheel_utils.detect_system_cuda_variant() |
| 927 | commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower() |
| 928 | if not commit or len(commit) != 40: |
| 929 | print( |
| 930 | f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}" |
| 931 | ", trying to fetch base commit in main branch" |
| 932 | ) |
| 933 | commit = precompiled_wheel_utils.get_base_commit_in_main_branch() |
| 934 | print(f"Using precompiled wheel commit {commit} with variant {variant}") |
| 935 | download_filename = None |
| 936 | try: |
| 937 | wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant( |
| 938 | commit, variant |
| 939 | ) |
| 940 | except Exception as e: |
| 941 | raise RuntimeError( |
| 942 | "Failed to fetch precompiled wheel metadata for CUDA " |
| 943 | f"variant {variant!r} at commit {commit}. The " |
| 944 | "root/default variant is not used as a fallback because " |
| 945 | "its CUDA compatibility with the selected variant cannot " |
| 946 | "be verified. Provide a compatible wheel with " |
| 947 | "VLLM_PRECOMPILED_WHEEL_LOCATION or disable " |
| 948 | "VLLM_USE_PRECOMPILED to build from source." |
| 949 | ) from e |
| 950 | # The metadata.json has the following format: |
| 951 | # see .buildkite/scripts/generate-nightly-index.py for details |
| 952 | """[{ |
| 953 | "package_name": "vllm", |
| 954 | "version": "0.11.2.dev278+gdbc3d9991", |
| 955 | "build_tag": null, |
| 956 | "python_tag": "cp38", |
| 957 | "abi_tag": "abi3", |
| 958 | "platform_tag": "manylinux1_x86_64", |
| 959 | "variant": null, |
| 960 | "filename": "vllm-0.11.2.dev278+gdbc3d9991-cp38-abi3-manylinux1_x86_64.whl", |
| 961 | "path": "../vllm-0.11.2.dev278%2Bgdbc3d9991-cp38-abi3-manylinux1_x86_64.whl" |
| 962 | }, |
| 963 | ...]""" |
| 964 | from urllib.parse import urljoin |
| 965 | |
| 966 | for wheel in wheels: |
| 967 | # TODO: maybe check more compatibility later? (python_tag, abi_tag, etc) |
| 968 | if wheel.get("package_name") == "vllm" and arch in wheel.get( |
| 969 | "platform_tag", "" |
| 970 | ): |
| 971 | print(f"Found precompiled wheel metadata: {wheel}") |
| 972 | if "path" not in wheel: |
| 973 | raise ValueError(f"Wheel metadata missing path: {wheel}") |
| 974 | wheel_url = urljoin(repo_url, wheel["path"]) |
| 975 | download_filename = wheel.get("filename") |
| 976 | print(f"Using precompiled wheel URL: {wheel_url}") |
| 977 | break |
| 978 | else: |
| 979 | raise ValueError( |
| 980 | f"No precompiled vllm wheel found for architecture {arch} " |
| 981 | f"from repo {repo_url}. All available wheels: {wheels}" |
| 982 | ) |
| 983 | |
| 984 | return wheel_url, download_filename |
| 985 | |
| 986 | @staticmethod |
| 987 | def extract_precompiled_and_patch_package( |
| 988 | wheel_url_or_path: str, |
| 989 | download_filename: str | None, |
| 990 | *, |
| 991 | extract_extensions: bool, |
| 992 | extract_rust_frontend: bool, |
| 993 | ) -> dict: |
| 994 | import tempfile |
| 995 | import zipfile |
| 996 | |
| 997 | temp_dir = None |
| 998 | try: |
| 999 | if not os.path.isfile(wheel_url_or_path): |
| 1000 | # use provided filename first, then derive from URL |
| 1001 | wheel_filename = download_filename or wheel_url_or_path.split("/")[-1] |
| 1002 | temp_dir = tempfile.mkdtemp(prefix="vllm-wheels") |
| 1003 | wheel_path = os.path.join(temp_dir, wheel_filename) |
| 1004 | print(f"Downloading wheel from {wheel_url_or_path} to {wheel_path}") |
| 1005 | from urllib.request import urlretrieve |
| 1006 | |
| 1007 | urlretrieve(wheel_url_or_path, filename=wheel_path) |
| 1008 | else: |
| 1009 | wheel_path = wheel_url_or_path |
| 1010 | print(f"Using existing wheel at {wheel_path}") |
| 1011 | |
| 1012 | package_data_patch = {} |
| 1013 | |
| 1014 | with zipfile.ZipFile(wheel_path) as wheel: |
| 1015 | exact_members = set() |
| 1016 | if extract_extensions: |
| 1017 | exact_members.update( |
| 1018 | { |
| 1019 | "vllm/_C.abi3.so", |
| 1020 | "vllm/_C_stable_libtorch.abi3.so", |
| 1021 | "vllm/_moe_C_stable_libtorch.abi3.so", |
| 1022 | "vllm/_qutlass_C.abi3.so", |
| 1023 | "vllm/_flashmla_C.abi3.so", |
| 1024 | "vllm/_flashmla_extension_C.abi3.so", |
| 1025 | "vllm/_flashkda_C.abi3.so", |
| 1026 | "vllm/_sparse_flashmla_C.abi3.so", |
| 1027 | "vllm/_deepselect_C.abi3.so", |
| 1028 | "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so", |
| 1029 | "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", |
| 1030 | "vllm/cumem_allocator.abi3.so", |
| 1031 | "vllm/spinloop.abi3.so", |
| 1032 | "vllm/fs_io_C.abi3.so", |
| 1033 | # ROCm-specific libraries |
| 1034 | "vllm/_rocm_C.abi3.so", |
| 1035 | } |
| 1036 | ) |
| 1037 | if extract_rust_frontend: |
| 1038 | exact_members.add("vllm/vllm-rs") |
| 1039 | |
| 1040 | flash_attn_regex = re.compile( |
| 1041 | r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" |
| 1042 | ) |
| 1043 | # __init__.py and flash_attn_interface.py are source-controlled |
| 1044 | # in vllm and should not be overwritten (matches cmake exclusions) |
| 1045 | flash_attn_files_to_skip = { |
| 1046 | "vllm/vllm_flash_attn/__init__.py", |
| 1047 | "vllm/vllm_flash_attn/flash_attn_interface.py", |
| 1048 | } |
| 1049 | triton_kernels_regex = re.compile( |
| 1050 | r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" |
| 1051 | ) |
| 1052 | flashmla_regex = re.compile( |
| 1053 | r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" |
| 1054 | ) |
| 1055 | # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.) |
| 1056 | deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*") |
| 1057 | fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*") |
| 1058 | tml_fa4_regex = re.compile(r"vllm/third_party/tml_fa4/.*") |
| 1059 | file_members = [] |
| 1060 | for member in wheel.filelist: |
| 1061 | if member.filename in exact_members: |
| 1062 | file_members.append(member) |
| 1063 | continue |
| 1064 | if ( |
| 1065 | extract_rust_frontend |
| 1066 | and PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX.match( |
| 1067 | member.filename |
| 1068 | ) |
| 1069 | ): |
| 1070 | file_members.append(member) |
| 1071 | continue |
| 1072 | |
| 1073 | if not extract_extensions: |
| 1074 | continue |
| 1075 | |
| 1076 | if ( |
| 1077 | ( |
| 1078 | flash_attn_regex.match(member.filename) |
| 1079 | and member.filename not in flash_attn_files_to_skip |
| 1080 | ) |
| 1081 | or triton_kernels_regex.match(member.filename) |
| 1082 | or flashmla_regex.match(member.filename) |
| 1083 | or deep_gemm_regex.match(member.filename) |
| 1084 | or tml_fa4_regex.match(member.filename) |
| 1085 | or fmha_sm100_regex.match(member.filename) |
| 1086 | ): |
| 1087 | file_members.append(member) |
| 1088 | |
| 1089 | for file in file_members: |
| 1090 | print(f"[extract] {file.filename}") |
| 1091 | target_path = os.path.join(".", file.filename) |
| 1092 | os.makedirs(os.path.dirname(target_path), exist_ok=True) |
| 1093 | with ( |
| 1094 | wheel.open(file.filename) as src, |
| 1095 | open(target_path, "wb") as dst, |
| 1096 | ): |
| 1097 | shutil.copyfileobj(src, dst) |
| 1098 | mode = file.external_attr >> 16 |
| 1099 | if mode: |
| 1100 | os.chmod(target_path, mode) |
| 1101 | |
| 1102 | pkg = os.path.dirname(file.filename).replace("/", ".") |
| 1103 | package_data_patch.setdefault(pkg, []).append( |
| 1104 | os.path.basename(file.filename) |
| 1105 | ) |
| 1106 | |
| 1107 | return package_data_patch |
| 1108 | finally: |
| 1109 | if temp_dir is not None: |
| 1110 | print(f"Removing temporary directory {temp_dir}") |
| 1111 | shutil.rmtree(temp_dir) |
| 1112 | |
| 1113 | @staticmethod |
| 1114 | def get_base_commit_in_main_branch() -> str: |
| 1115 | try: |
| 1116 | # Get the latest commit hash of the upstream main branch. |
| 1117 | curl_cmd = [ |
| 1118 | "curl", |
| 1119 | "-s", |
| 1120 | "https://api.github.com/repos/vllm-project/vllm/commits/main", |
| 1121 | ] |
| 1122 | github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN")) |
| 1123 | if github_token: |
| 1124 | curl_cmd += [ |
| 1125 | "-H", |
| 1126 | f"Authorization: token {github_token}", |
| 1127 | ] |
| 1128 | resp_json = subprocess.check_output(curl_cmd).decode("utf-8") |
| 1129 | upstream_main_commit = json.loads(resp_json)["sha"] |
| 1130 | print(f"Upstream main branch latest commit: {upstream_main_commit}") |
| 1131 | |
| 1132 | # In Docker build context, .git may be immutable or missing. |
| 1133 | if envs.VLLM_DOCKER_BUILD_CONTEXT: |
| 1134 | return upstream_main_commit |
| 1135 | |
| 1136 | # Check if the upstream_main_commit exists in the local repo |
| 1137 | try: |
| 1138 | subprocess.check_output( |
| 1139 | ["git", "cat-file", "-e", f"{upstream_main_commit}"] |
| 1140 | ) |
| 1141 | except subprocess.CalledProcessError: |
| 1142 | # If not present, fetch it from the remote repository. |
| 1143 | # Note that this does not update any local branches, |
| 1144 | # but ensures that this commit ref and its history are |
| 1145 | # available in our local repo. |
| 1146 | subprocess.check_call( |
| 1147 | ["git", "fetch", "https://github.com/vllm-project/vllm", "main"] |
| 1148 | ) |
| 1149 | |
| 1150 | # Then get the commit hash of the current branch that is the same as |
| 1151 | # the upstream main commit. |
| 1152 | current_branch = ( |
| 1153 | subprocess.check_output(["git", "branch", "--show-current"]) |
| 1154 | .decode("utf-8") |
| 1155 | .strip() |
| 1156 | ) |
| 1157 | |
| 1158 | base_commit = ( |
| 1159 | subprocess.check_output( |
| 1160 | ["git", "merge-base", f"{upstream_main_commit}", current_branch] |
| 1161 | ) |
| 1162 | .decode("utf-8") |
| 1163 | .strip() |
| 1164 | ) |
| 1165 | return base_commit |
| 1166 | except ValueError as err: |
| 1167 | raise ValueError(err) from None |
| 1168 | except Exception as err: |
| 1169 | logger.warning( |
| 1170 | "Failed to get the base commit in the main branch. " |
| 1171 | "Using the nightly wheel. The libraries in this " |
| 1172 | "wheel may not be compatible with your dev branch: %s", |
| 1173 | err, |
| 1174 | ) |
| 1175 | return "nightly" |
| 1176 | |
| 1177 | |
| 1178 | def _no_device() -> bool: |
| 1179 | return VLLM_TARGET_DEVICE == "empty" |
| 1180 | |
| 1181 | |
| 1182 | def _is_cuda() -> bool: |
| 1183 | has_cuda = torch.version.cuda is not None |
| 1184 | return VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu() |
| 1185 | |
| 1186 | |
| 1187 | def _is_hip() -> bool: |
| 1188 | return ( |
| 1189 | VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm" |
| 1190 | ) and torch.version.hip is not None |
| 1191 | |
| 1192 | |
| 1193 | def _is_tpu() -> bool: |
| 1194 | return VLLM_TARGET_DEVICE == "tpu" |
| 1195 | |
| 1196 | |
| 1197 | def _is_cpu() -> bool: |
| 1198 | return VLLM_TARGET_DEVICE == "cpu" |
| 1199 | |
| 1200 | |
| 1201 | def _is_xpu() -> bool: |
| 1202 | return VLLM_TARGET_DEVICE == "xpu" |
| 1203 | |
| 1204 | |
| 1205 | def _build_custom_ops() -> bool: |
| 1206 | return _is_cuda() or _is_hip() |
| 1207 | |
| 1208 | |
| 1209 | def get_rocm_version(): |
| 1210 | # Get the Rocm version from the ROCM_HOME/bin/librocm-core.so |
| 1211 | # see https://github.com/ROCm/rocm-core/blob/d11f5c20d500f729c393680a01fa902ebf92094b/rocm_version.cpp#L21 |
| 1212 | try: |
| 1213 | if ROCM_HOME is None: |
| 1214 | return None |
| 1215 | librocm_core_file = Path(ROCM_HOME) / "lib" / "librocm-core.so" |
| 1216 | if not librocm_core_file.is_file(): |
| 1217 | return None |
| 1218 | librocm_core = ctypes.CDLL(librocm_core_file) |
| 1219 | VerErrors = ctypes.c_uint32 |
| 1220 | get_rocm_core_version = librocm_core.getROCmVersion |
| 1221 | get_rocm_core_version.restype = VerErrors |
| 1222 | get_rocm_core_version.argtypes = [ |
| 1223 | ctypes.POINTER(ctypes.c_uint32), |
| 1224 | ctypes.POINTER(ctypes.c_uint32), |
| 1225 | ctypes.POINTER(ctypes.c_uint32), |
| 1226 | ] |
| 1227 | major = ctypes.c_uint32() |
| 1228 | minor = ctypes.c_uint32() |
| 1229 | patch = ctypes.c_uint32() |
| 1230 | |
| 1231 | if ( |
| 1232 | get_rocm_core_version( |
| 1233 | ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch) |
| 1234 | ) |
| 1235 | == 0 |
| 1236 | ): |
| 1237 | return f"{major.value}.{minor.value}.{patch.value}" |
| 1238 | return None |
| 1239 | except Exception: |
| 1240 | return None |
| 1241 | |
| 1242 | |
| 1243 | def get_nvcc_cuda_version() -> Version: |
| 1244 | """Get the CUDA version from nvcc. |
| 1245 | |
| 1246 | Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py |
| 1247 | """ |
| 1248 | assert CUDA_HOME is not None, "CUDA_HOME is not set" |
| 1249 | nvcc_output = subprocess.check_output( |
| 1250 | [CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True |
| 1251 | ) |
| 1252 | output = nvcc_output.split() |
| 1253 | release_idx = output.index("release") + 1 |
| 1254 | nvcc_cuda_version = parse(output[release_idx].split(",")[0]) |
| 1255 | return nvcc_cuda_version |
| 1256 | |
| 1257 | |
| 1258 | def get_vllm_version() -> str: |
| 1259 | # Allow overriding the version. This is useful to build platform-specific |
| 1260 | # wheels (e.g. CPU, TPU) without modifying the source. |
| 1261 | if env_version := os.getenv("VLLM_VERSION_OVERRIDE"): |
| 1262 | print(f"Overriding VLLM version with {env_version} from VLLM_VERSION_OVERRIDE") |
| 1263 | os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version |
| 1264 | return get_version( |
| 1265 | write_to="vllm/_version.py", |
| 1266 | git_describe_command=VLLM_GIT_DESCRIBE_COMMAND, |
| 1267 | ) |
| 1268 | |
| 1269 | version = get_version( |
| 1270 | write_to="vllm/_version.py", |
| 1271 | git_describe_command=VLLM_GIT_DESCRIBE_COMMAND, |
| 1272 | ) |
| 1273 | sep = "+" if "+" not in version else "." # dev versions might contain + |
| 1274 | |
| 1275 | if not envs.VLLM_SKIP_VERSION_SUFFIX: |
| 1276 | if _no_device(): |
| 1277 | if envs.VLLM_TARGET_DEVICE == "empty": |
| 1278 | version += f"{sep}empty" |
| 1279 | elif _is_cuda(): |
| 1280 | if ( |
| 1281 | USE_PRECOMPILED_EXTENSIONS |
| 1282 | and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX |
| 1283 | ): |
| 1284 | version += f"{sep}precompiled" |
| 1285 | else: |
| 1286 | cuda_version = str(get_nvcc_cuda_version()) |
| 1287 | if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: |
| 1288 | cuda_version_str = cuda_version.replace(".", "")[:3] |
| 1289 | # skip this for source tarball, required for pypi |
| 1290 | if "sdist" not in sys.argv: |
| 1291 | version += f"{sep}cu{cuda_version_str}" |
| 1292 | elif _is_hip(): |
| 1293 | # Get the Rocm Version |
| 1294 | rocm_version = get_rocm_version() or torch.version.hip |
| 1295 | if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: |
| 1296 | version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" |
| 1297 | elif _is_tpu(): |
| 1298 | version += f"{sep}tpu" |
| 1299 | elif _is_cpu(): |
| 1300 | # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), |
| 1301 | # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. |
| 1302 | if VLLM_TARGET_DEVICE == "cpu": |
| 1303 | version += f"{sep}cpu" |
| 1304 | elif _is_xpu(): |
| 1305 | version += f"{sep}xpu" |
| 1306 | else: |
| 1307 | raise RuntimeError("Unknown runtime environment") |
| 1308 | |
| 1309 | return version |
| 1310 | |
| 1311 | |
| 1312 | def get_requirements() -> list[str]: |
| 1313 | """Get Python package dependencies from requirements.txt.""" |
| 1314 | requirements_dir = ROOT_DIR / "requirements" |
| 1315 | |
| 1316 | def _read_requirements(filename: str) -> list[str]: |
| 1317 | with open(requirements_dir / filename) as f: |
| 1318 | requirements = f.read().strip().split("\n") |
| 1319 | resolved_requirements = [] |
| 1320 | for line in requirements: |
| 1321 | if line.startswith("-r "): |
| 1322 | resolved_requirements += _read_requirements(line.split()[1]) |
| 1323 | elif ( |
| 1324 | not line.startswith("--") |
| 1325 | and not line.startswith("#") |
| 1326 | and line.strip() != "" |
| 1327 | ): |
| 1328 | resolved_requirements.append(line) |
| 1329 | return resolved_requirements |
| 1330 | |
| 1331 | if _no_device(): |
| 1332 | requirements = _read_requirements("common.txt") |
| 1333 | elif _is_cuda(): |
| 1334 | requirements = _read_requirements("cuda.txt") |
| 1335 | cuda_major, cuda_minor = torch.version.cuda.split(".") |
| 1336 | modified_requirements = [] |
| 1337 | for req in requirements: |
| 1338 | if "vllm-flash-attn" in req and cuda_major != "12": |
| 1339 | # vllm-flash-attn is built only for CUDA 12.x. |
| 1340 | # Skip for other versions. |
| 1341 | continue |
| 1342 | if "flashinfer-cubin" in req: |
| 1343 | # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so |
| 1344 | # it cannot be a wheel dependency; flashinfer falls back to |
| 1345 | # fetching cubins at runtime when the package is absent. |
| 1346 | continue |
| 1347 | if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": |
| 1348 | # [cu13] extra is the default; strip it on CUDA 12 builds. |
| 1349 | req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl") |
| 1350 | if "humming-kernels[cu13]" in req and cuda_major == "12": |
| 1351 | req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]") |
| 1352 | modified_requirements.append(req) |
| 1353 | requirements = modified_requirements |
| 1354 | elif _is_hip(): |
| 1355 | requirements = _read_requirements("rocm.txt") |
| 1356 | elif _is_tpu(): |
| 1357 | requirements = _read_requirements("tpu.txt") |
| 1358 | elif _is_cpu(): |
| 1359 | requirements = _read_requirements("cpu.txt") |
| 1360 | elif _is_xpu(): |
| 1361 | requirements = _read_requirements("xpu.txt") |
| 1362 | else: |
| 1363 | raise ValueError("Unsupported platform, please use CUDA, ROCm, or CPU.") |
| 1364 | return requirements |
| 1365 | |
| 1366 | |
| 1367 | ext_modules = [] |
| 1368 | |
| 1369 | if _is_cuda() or _is_hip(): |
| 1370 | ext_modules.append(CMakeExtension(name="vllm.cumem_allocator")) |
| 1371 | # Optional since this doesn't get built (produce an .so file). This is just |
| 1372 | # copying the relevant .py files from the source repository. |
| 1373 | ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True)) |
| 1374 | |
| 1375 | if not _is_xpu() and sys.version_info >= (3, 11): |
| 1376 | ext_modules.append(CMakeExtension(name="vllm.spinloop")) |
| 1377 | ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) |
| 1378 | |
| 1379 | if _is_hip(): |
| 1380 | ext_modules.append(CMakeExtension(name="vllm._rocm_C")) |
| 1381 | |
| 1382 | if _is_cuda(): |
| 1383 | ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) |
| 1384 | if USE_PRECOMPILED_EXTENSIONS or ( |
| 1385 | CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") |
| 1386 | ): |
| 1387 | # FA3 requires CUDA 12.3 or later |
| 1388 | ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C")) |
| 1389 | # FA4 CuteDSL - Python-only component for FA4's cute DSL support |
| 1390 | # Optional since this doesn't produce a .so file, just copies Python files |
| 1391 | ext_modules.append( |
| 1392 | CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True) |
| 1393 | ) |
| 1394 | if USE_PRECOMPILED_EXTENSIONS or ( |
| 1395 | CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9") |
| 1396 | ): |
| 1397 | # FlashMLA requires CUDA 12.9 or later |
| 1398 | # Optional since this doesn't get built (produce an .so file) when |
| 1399 | # not targeting a hopper system |
| 1400 | ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True)) |
| 1401 | ext_modules.append( |
| 1402 | CMakeExtension(name="vllm._flashmla_extension_C", optional=True) |
| 1403 | ) |
| 1404 | # DeepSelect requires CUDA 12.9 or later (SM100a/SM103a only) |
| 1405 | # Optional since it won't build on unsupported architectures |
| 1406 | ext_modules.append(CMakeExtension(name="vllm._deepselect_C", optional=True)) |
| 1407 | if USE_PRECOMPILED_EXTENSIONS or ( |
| 1408 | CUDA_HOME and get_nvcc_cuda_version() >= Version("12.0") |
| 1409 | ): |
| 1410 | ext_modules.append(CMakeExtension(name="vllm._flashkda_C", optional=True)) |
| 1411 | if envs.VLLM_USE_PRECOMPILED or ( |
| 1412 | CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") |
| 1413 | ): |
| 1414 | # DeepGEMM requires CUDA 12.3+ (SM90/SM100) |
| 1415 | # Optional since it won't build on unsupported architectures |
| 1416 | ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True)) |
| 1417 | ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True)) |
| 1418 | # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party. |
| 1419 | ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True)) |
| 1420 | # tml-fa4 is copied into an isolated vllm.third_party package. |
| 1421 | ext_modules.append(CMakeExtension(name="vllm.tml_fa4", optional=True)) |
| 1422 | |
| 1423 | if _is_cpu(): |
| 1424 | import platform |
| 1425 | |
| 1426 | if platform.machine() in ("x86_64", "AMD64"): |
| 1427 | ext_modules.append(CMakeExtension(name="vllm._C")) |
| 1428 | ext_modules.append(CMakeExtension(name="vllm._C_AVX512")) |
| 1429 | ext_modules.append(CMakeExtension(name="vllm._C_AVX2")) |
| 1430 | else: |
| 1431 | ext_modules.append(CMakeExtension(name="vllm._C")) |
| 1432 | |
| 1433 | if _build_custom_ops(): |
| 1434 | if _is_hip(): |
| 1435 | ext_modules.append(CMakeExtension(name="vllm._C")) |
| 1436 | if _is_cuda() or _is_hip(): |
| 1437 | ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch")) |
| 1438 | ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch")) |
| 1439 | |
| 1440 | package_data = { |
| 1441 | "vllm": [ |
| 1442 | "py.typed", |
| 1443 | "libs/*.so*", |
| 1444 | "model_executor/layers/fused_moe/configs/*.json", |
| 1445 | "model_executor/layers/quantization/utils/configs/*.json", |
| 1446 | "entrypoints/serve/instrumentator/static/*.js", |
| 1447 | "entrypoints/serve/instrumentator/static/*.css", |
| 1448 | "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp", |
| 1449 | # Built-in multimodal chat template fallbacks (registry.py) |
| 1450 | "transformers_utils/chat_templates/*.jinja", |
| 1451 | "third_party/flash_linear_attention/LICENSE", |
| 1452 | # DeepGEMM JIT include headers (vendored via cmake) |
| 1453 | "third_party/deep_gemm/include/**/*.cuh", |
| 1454 | "third_party/deep_gemm/include/**/*.h", |
| 1455 | "third_party/deep_gemm/include/**/*.hpp", |
| 1456 | # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake) |
| 1457 | "third_party/fmha_sm100/csrc/**/*.cu", |
| 1458 | "third_party/fmha_sm100/csrc/**/*.h", |
| 1459 | "third_party/fmha_sm100/csrc/**/*.jinja", |
| 1460 | "third_party/fmha_sm100/csrc/**/*.cu.jinja", |
| 1461 | "third_party/fmha_sm100/cute/**/*.cu", |
| 1462 | "third_party/fmha_sm100/cutlass/include/**/*.h", |
| 1463 | "third_party/fmha_sm100/cutlass/include/**/*.hpp", |
| 1464 | "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h", |
| 1465 | "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp", |
| 1466 | ] |
| 1467 | } |
| 1468 | |
| 1469 | |
| 1470 | def add_vllm_package_data(filename: str) -> None: |
| 1471 | vllm_files = package_data.setdefault("vllm", []) |
| 1472 | if filename not in vllm_files: |
| 1473 | vllm_files.append(filename) |
| 1474 | |
| 1475 | |
| 1476 | # PEP 517 invokes setup.py for metadata before invoking the actual build. |
| 1477 | if USE_PRECOMPILED_RUST_FRONTEND and not is_metadata_only_build(): |
| 1478 | wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url() |
| 1479 | patch = precompiled_wheel_utils.extract_precompiled_and_patch_package( |
| 1480 | wheel_url, |
| 1481 | download_filename, |
| 1482 | extract_extensions=USE_PRECOMPILED_EXTENSIONS, |
| 1483 | extract_rust_frontend=True, |
| 1484 | ) |
| 1485 | for pkg, files in patch.items(): |
| 1486 | package_data.setdefault(pkg, []).extend(files) |
| 1487 | |
| 1488 | # If the rust frontend binary is already present in the source tree (e.g., |
| 1489 | # pre-built in a separate Docker build stage), ship it as-is. |
| 1490 | if PRECOMPILED_RUST_FRONTEND_PATH.exists(): |
| 1491 | add_vllm_package_data("vllm-rs") |
| 1492 | for rust_extension_path in get_precompiled_rust_extension_paths(): |
| 1493 | add_vllm_package_data(rust_extension_path.name) |
| 1494 | |
| 1495 | if _no_device(): |
| 1496 | ext_modules = [] |
| 1497 | |
| 1498 | if not ext_modules: |
| 1499 | cmdclass = {} |
| 1500 | else: |
| 1501 | cmdclass = { |
| 1502 | "build_ext": precompiled_build_ext |
| 1503 | if USE_PRECOMPILED_EXTENSIONS |
| 1504 | else cmake_build_ext, |
| 1505 | } |
| 1506 | if ( |
| 1507 | USE_PRECOMPILED_RUST_FRONTEND |
| 1508 | or PRECOMPILED_RUST_FRONTEND_PATH.exists() |
| 1509 | or has_precompiled_rust_extensions() |
| 1510 | ): |
| 1511 | cmdclass["build_rust"] = precompiled_build_rust |
| 1512 | |
| 1513 | # Resolve the Python version first because get_vllm_version() may set |
| 1514 | # SETUPTOOLS_SCM_PRETEND_VERSION, which the Rust version should inherit. |
| 1515 | vllm_version = get_vllm_version() |
| 1516 | rust_build.prepare_build_environment() |
| 1517 | |
| 1518 | # Rust artifacts, built via setuptools-rust and installed into the package |
| 1519 | # directory alongside the Python modules. |
| 1520 | rust_extensions = rust_build.rust_extensions( |
| 1521 | optional=not should_require_rust_frontend() |
| 1522 | ) |
| 1523 | |
| 1524 | setup( |
| 1525 | # static metadata should rather go in pyproject.toml |
| 1526 | version=vllm_version, |
| 1527 | ext_modules=ext_modules, |
| 1528 | rust_extensions=rust_extensions, |
| 1529 | install_requires=get_requirements(), |
| 1530 | extras_require={ |
| 1531 | # AMD Zen CPU optimizations via zentorch |
| 1532 | "zen": ["zentorch==2.13.0.1"], |
| 1533 | "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], |
| 1534 | "tensorizer": ["tensorizer==2.10.1"], |
| 1535 | "fastsafetensors": ["fastsafetensors >= 0.3.3"], |
| 1536 | "instanttensor": ["instanttensor >= 0.1.9"], |
| 1537 | "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"], |
| 1538 | "audio": [ |
| 1539 | "av", |
| 1540 | "scipy", |
| 1541 | "soundfile", |
| 1542 | "soxr", |
| 1543 | "mistral_common[audio]", |
| 1544 | ], # Required for audio processing |
| 1545 | "video": [], # Kept for backwards compatibility |
| 1546 | # NVIDIA DeepStream (NVDEC) GPU video-decode backend. Linux x86-64 |
| 1547 | # only; also needs system GStreamer + libv4l (see docs). |
| 1548 | "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], |
| 1549 | "flashinfer": [], # Kept for backwards compatibility |
| 1550 | "b12x": ["b12x==1.3.0"], |
| 1551 | # Optional deps for Helion kernel development |
| 1552 | # NOTE: When updating helion version, also update CI files: |
| 1553 | # - .buildkite/test_areas/kernels.yaml |
| 1554 | # - .buildkite/test-amd.yaml |
| 1555 | "helion": ["helion==1.4.0"], |
| 1556 | # Optional deps for gRPC server (vllm serve --grpc) |
| 1557 | "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"], |
| 1558 | # Optional deps for OpenTelemetry tracing |
| 1559 | "otel": [ |
| 1560 | "opentelemetry-sdk>=1.26.0", |
| 1561 | "opentelemetry-api>=1.26.0", |
| 1562 | "opentelemetry-exporter-otlp>=1.26.0", |
| 1563 | "opentelemetry-semantic-conventions-ai>=0.4.1", |
| 1564 | ], |
| 1565 | # extra quantization plugin |
| 1566 | "extra-quant": ["vllm-gguf-plugin>=0.0.2"], |
| 1567 | }, |
| 1568 | cmdclass=cmdclass, |
| 1569 | package_data=package_data, |
| 1570 | ) |