setup.py 34.1 KB
Newer Older
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
9
import os
import re
10
import shutil
11
import subprocess
bnellnm's avatar
bnellnm committed
12
import sys
13
import sysconfig
14
from pathlib import Path
15
from shutil import which
16

Woosuk Kwon's avatar
Woosuk Kwon committed
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_scm import get_version
22
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
23

zhuwenwen's avatar
zhuwenwen committed
24
25
26
27
from typing import Optional, Union
import subprocess
from pathlib import Path

zhuwenwen's avatar
zhuwenwen committed
28
pwd = os.path.dirname(os.path.abspath(__file__))
gaoqiong's avatar
gaoqiong committed
29
add_git_version = False
30
31
if int(os.environ.get('ADD_GIT_VERSION', '0')) == 1:
    add_git_version = True
zhuwenwen's avatar
zhuwenwen committed
32
33
34
35
    
skip_vllm_build = False
if int(os.environ.get('SKIP_VLLM_BUILD', '0')) == 1:
    skip_vllm_build = True
36
37
38
39
40
41
42
43
44

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


45
ROOT_DIR = Path(__file__).parent
46
logger = logging.getLogger(__name__)
47
48
49

# cannot import envs directly because it depends on vllm,
#  which is not installed yet
50
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
51
52

VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
53

54
if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
55
    logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
56
    VLLM_TARGET_DEVICE = "cpu"
57
elif not (sys.platform.startswith("linux") or sys.platform.startswith("darwin")):
58
59
    logger.warning(
        "vLLM only supports Linux platform (including WSL) and MacOS."
60
        "Building on %s, "
61
62
63
        "so vLLM may not be able to run correctly",
        sys.platform,
    )
64
    VLLM_TARGET_DEVICE = "empty"
65
66
67
68
69
70
elif (
    sys.platform.startswith("linux")
    and torch.version.cuda is None
    and os.getenv("VLLM_TARGET_DEVICE") is None
    and torch.version.hip is None
):
71
    # if cuda or hip is not available and VLLM_TARGET_DEVICE is not set,
72
73
    # fallback to cpu
    VLLM_TARGET_DEVICE = "cpu"
74

bnellnm's avatar
bnellnm committed
75
76

def is_sccache_available() -> bool:
77
78
79
    return which("sccache") is not None and not bool(
        int(os.getenv("VLLM_DISABLE_SCCACHE", "0"))
    )
bnellnm's avatar
bnellnm committed
80
81
82
83
84
85
86
87
88
89


def is_ccache_available() -> bool:
    return which("ccache") is not None


def is_ninja_available() -> bool:
    return which("ninja") is not None


90
91
def is_freethreaded():
    return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
92
93


bnellnm's avatar
bnellnm committed
94
class CMakeExtension(Extension):
95
    def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
96
        super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
bnellnm's avatar
bnellnm committed
97
98
99
100
101
        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.
102
    did_config: dict[str, bool] = {}
bnellnm's avatar
bnellnm committed
103
104
105
106
107

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
108
109
        # `num_jobs` is either the value of the MAX_JOBS environment variable
        # (if defined) or the number of CPUs available.
110
        num_jobs = envs.MAX_JOBS
111
112
        if num_jobs is not None:
            num_jobs = int(num_jobs)
113
            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
114
115
116
117
118
119
120
        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()
bnellnm's avatar
bnellnm committed
121

122
        nvcc_threads = None
123
124
125
126
127
        if _is_cuda() and get_nvcc_cuda_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.
128
            nvcc_threads = envs.NVCC_THREADS
129
130
            if nvcc_threads is not None:
                nvcc_threads = int(nvcc_threads)
131
                logger.info(
132
133
                    "Using NVCC_THREADS=%d as the number of nvcc threads.", nvcc_threads
                )
134
135
136
            else:
                nvcc_threads = 1
            num_jobs = max(1, num_jobs // nvcc_threads)
bnellnm's avatar
bnellnm committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

        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"
154
        cfg = envs.CMAKE_BUILD_TYPE or default_cfg
bnellnm's avatar
bnellnm committed
155
156

        cmake_args = [
157
158
            "-DCMAKE_BUILD_TYPE={}".format(cfg),
            "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
bnellnm's avatar
bnellnm committed
159
160
        ]

161
        verbose = envs.VERBOSE
bnellnm's avatar
bnellnm committed
162
        if verbose:
163
            cmake_args += ["-DCMAKE_VERBOSE_MAKEFILE=ON"]
bnellnm's avatar
bnellnm committed
164
165
166

        if is_sccache_available():
            cmake_args += [
167
168
169
170
                "-DCMAKE_C_COMPILER_LAUNCHER=sccache",
                "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache",
                "-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache",
                "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache",
bnellnm's avatar
bnellnm committed
171
172
173
            ]
        elif is_ccache_available():
            cmake_args += [
174
175
176
177
                "-DCMAKE_C_COMPILER_LAUNCHER=ccache",
                "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
                "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache",
                "-DCMAKE_HIP_COMPILER_LAUNCHER=ccache",
bnellnm's avatar
bnellnm committed
178
179
180
181
            ]

        # Pass the python executable to cmake so it can find an exact
        # match.
182
        cmake_args += ["-DVLLM_PYTHON_EXECUTABLE={}".format(sys.executable)]
bnellnm's avatar
bnellnm committed
183

184
185
        # Pass the python path to cmake so it can reuse the build dependencies
        # on subsequent calls to python.
186
        cmake_args += ["-DVLLM_PYTHON_PATH={}".format(":".join(sys.path))]
187

188
189
190
191
192
193
        # 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)
194
        cmake_args += ["-DFETCHCONTENT_BASE_DIR={}".format(fc_base_dir)]
195

bnellnm's avatar
bnellnm committed
196
197
198
199
200
201
        #
        # Setup parallelism and build tool
        #
        num_jobs, nvcc_threads = self.compute_num_jobs()

        if nvcc_threads:
202
            cmake_args += ["-DNVCC_THREADS={}".format(nvcc_threads)]
bnellnm's avatar
bnellnm committed
203
204

        if is_ninja_available():
205
            build_tool = ["-G", "Ninja"]
bnellnm's avatar
bnellnm committed
206
            cmake_args += [
207
208
                "-DCMAKE_JOB_POOL_COMPILE:STRING=compile",
                "-DCMAKE_JOB_POOLS:STRING=compile={}".format(num_jobs),
bnellnm's avatar
bnellnm committed
209
210
211
212
            ]
        else:
            # Default build tool to whatever cmake picks.
            build_tool = []
213
214
        # Make sure we use the nvcc from CUDA_HOME
        if _is_cuda():
215
            cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"]
216
217
        elif _is_hip():
            cmake_args += [f"-DROCM_PATH={ROCM_HOME}"]
218
219
220
221
222

        other_cmake_args = os.environ.get("CMAKE_ARGS")
        if other_cmake_args:
            cmake_args += other_cmake_args.split()

bnellnm's avatar
bnellnm committed
223
        subprocess.check_call(
224
225
226
            ["cmake", ext.cmake_lists_dir, *build_tool, *cmake_args],
            cwd=self.build_temp,
        )
bnellnm's avatar
bnellnm committed
227
228
229
230

    def build_extensions(self) -> None:
        # Ensure that CMake is present and working
        try:
231
            subprocess.check_output(["cmake", "--version"])
bnellnm's avatar
bnellnm committed
232
        except OSError as e:
233
            raise RuntimeError("Cannot find CMake executable") from e
bnellnm's avatar
bnellnm committed
234
235
236
237
238

        # Create build directory if it does not exist.
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

239
        targets = []
240
241
242
243

        def target_name(s: str) -> str:
            return s.removeprefix("vllm.").removeprefix("vllm_flash_attn.")

bnellnm's avatar
bnellnm committed
244
245
246
        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
247
            targets.append(target_name(ext.name))
bnellnm's avatar
bnellnm committed
248

249
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
250

251
252
253
254
255
256
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
bnellnm's avatar
bnellnm committed
257

258
        subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)
bnellnm's avatar
bnellnm committed
259

260
261
262
263
264
265
266
267
268
269
270
271
        # 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
272
            for _ in range(ext.name.count(".")):
273
274
275
276
                prefix = prefix.parent

            # prefix here should actually be the same for all components
            install_args = [
277
278
279
280
281
282
283
                "cmake",
                "--install",
                ".",
                "--prefix",
                prefix,
                "--component",
                target_name(ext.name),
284
285
286
            ]
            subprocess.check_call(install_args, cwd=self.build_temp)

287
288
289
290
    def run(self):
        # First, run the standard build_ext command to compile the extensions
        super().run()

291
        # copy vllm/vllm_flash_attn/**/*.py from self.build_lib to current
292
293
        # directory so that they can be included in the editable build
        import glob
294
295
296
297
298

        files = glob.glob(
            os.path.join(self.build_lib, "vllm", "vllm_flash_attn", "**", "*.py"),
            recursive=True,
        )
299
        for file in files:
300
301
302
            dst_file = os.path.join(
                "vllm/vllm_flash_attn", file.split("vllm/vllm_flash_attn/")[-1]
            )
303
            print(f"Copying {file} to {dst_file}")
304
            os.makedirs(os.path.dirname(dst_file), exist_ok=True)
305
306
            self.copy_file(file, dst_file)

307
        if _is_cuda():
308
309
310
311
312
313
314
315
316
317
318
319
320
            # 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,
            )

321

322
323
324
325
class precompiled_build_ext(build_ext):
    """Disables extension building when using precompiled binaries."""

    def run(self) -> None:
326
        return
327
328
329
330
331
332
333

    def build_extensions(self) -> None:
        print("Skipping build_ext: using precompiled extensions.")
        return


class precompiled_wheel_utils:
334
335
    """Extracts libraries and other files from an existing wheel."""

336
    @staticmethod
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
    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 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)
361
        2. user-specified variant (VLLM_PRECOMPILED_WHEEL_VARIANT) from nightly repo
362
        3. the variant corresponding to VLLM_MAIN_CUDA_VERSION from nightly repo
363
364
365
366
367
        4. 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.
368
369
370
371
372
373
374
375
376
377
378
379
        """
        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:
            import platform

            arch = platform.machine()
            # try to fetch the wheel metadata from the nightly wheel repo
            main_variant = "cu" + envs.VLLM_MAIN_CUDA_VERSION.replace(".", "")
            variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", main_variant)
380
381
382
383
384
385
386
            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()
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
            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
    ) -> dict:
450
451
452
453
454
455
        import tempfile
        import zipfile

        temp_dir = None
        try:
            if not os.path.isfile(wheel_url_or_path):
456
457
                # use provided filename first, then derive from URL
                wheel_filename = download_filename or wheel_url_or_path.split("/")[-1]
458
459
                temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
                wheel_path = os.path.join(temp_dir, wheel_filename)
460
                print(f"Downloading wheel from {wheel_url_or_path} to {wheel_path}")
461
                from urllib.request import urlretrieve
462

463
464
465
466
467
468
469
470
471
472
473
474
                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:
                files_to_copy = [
                    "vllm/_C.abi3.so",
                    "vllm/_moe_C.abi3.so",
                    "vllm/_flashmla_C.abi3.so",
475
476
                    "vllm/_flashmla_extension_C.abi3.so",
                    "vllm/_sparse_flashmla_C.abi3.so",
477
478
479
480
481
                    "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                    "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
                    "vllm/cumem_allocator.abi3.so",
                ]

482
                flash_attn_regex = re.compile(
483
484
                    r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                )
485
486
487
                triton_kernels_regex = re.compile(
                    r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                )
488
                file_members = list(
489
490
                    filter(lambda x: x.filename in files_to_copy, wheel.filelist)
                )
491
                file_members += list(
492
493
494
495
496
497
                    filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
                )
                file_members += list(
                    filter(
                        lambda x: triton_kernels_regex.match(x.filename), wheel.filelist
                    )
498
                )
499
500
501
502
503

                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)
504
505
506
507
                    with (
                        wheel.open(file.filename) as src,
                        open(target_path, "wb") as dst,
                    ):
508
509
510
511
                        shutil.copyfileobj(src, dst)

                    pkg = os.path.dirname(file.filename).replace("/", ".")
                    package_data_patch.setdefault(pkg, []).append(
512
513
                        os.path.basename(file.filename)
                    )
514
515
516
517
518
519
520
521
522

            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:
523
        try:
524
            # Get the latest commit hash of the upstream main branch.
525
526
527
528
529
530
531
            resp_json = subprocess.check_output(
                [
                    "curl",
                    "-s",
                    "https://api.github.com/repos/vllm-project/vllm/commits/main",
                ]
            ).decode("utf-8")
532
            upstream_main_commit = json.loads(resp_json)["sha"]
533
            print(f"Upstream main branch latest commit: {upstream_main_commit}")
534

535
536
537
538
            # In Docker build context, .git may be immutable or missing.
            if envs.VLLM_DOCKER_BUILD_CONTEXT:
                return upstream_main_commit

539
540
541
            # Check if the upstream_main_commit exists in the local repo
            try:
                subprocess.check_output(
542
543
                    ["git", "cat-file", "-e", f"{upstream_main_commit}"]
                )
544
545
546
547
548
            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.
549
550
551
                subprocess.check_call(
                    ["git", "fetch", "https://github.com/vllm-project/vllm", "main"]
                )
552
553
554

            # Then get the commit hash of the current branch that is the same as
            # the upstream main commit.
555
556
557
558
559
            current_branch = (
                subprocess.check_output(["git", "branch", "--show-current"])
                .decode("utf-8")
                .strip()
            )
560

561
562
563
564
565
566
567
            base_commit = (
                subprocess.check_output(
                    ["git", "merge-base", f"{upstream_main_commit}", current_branch]
                )
                .decode("utf-8")
                .strip()
            )
568
            return base_commit
569
570
        except ValueError as err:
            raise ValueError(err) from None
571
572
573
574
        except Exception as err:
            logger.warning(
                "Failed to get the base commit in the main branch. "
                "Using the nightly wheel. The libraries in this "
575
576
577
                "wheel may not be compatible with your dev branch: %s",
                err,
            )
578
            return "nightly"
579
580


581
582
583
584
def _no_device() -> bool:
    return VLLM_TARGET_DEVICE == "empty"


585
def _is_cuda() -> bool:
586
    has_cuda = torch.version.cuda is not None
587
    return VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu()
588
589
590


def _is_hip() -> bool:
591
592
593
    return (
        VLLM_TARGET_DEVICE == "cuda" or VLLM_TARGET_DEVICE == "rocm"
    ) and torch.version.hip is not None
594
595


596
597
598
599
def _is_tpu() -> bool:
    return VLLM_TARGET_DEVICE == "tpu"


600
601
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"
602
603


604
605
606
607
def _is_xpu() -> bool:
    return VLLM_TARGET_DEVICE == "xpu"


608
def _build_custom_ops() -> bool:
zhuwenwen's avatar
zhuwenwen committed
609
610
    if not skip_vllm_build:
        return _is_cuda() or _is_hip() or _is_cpu()
611
612


613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
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:
        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()
632

633
634
635
636
637
638
        if (
            get_rocm_core_version(
                ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
            )
            == 0
        ):
639
            return f"{major.value}.{minor.value}.{patch.value}"
640
        return None
641
    except Exception:
642
        return None
Woosuk Kwon's avatar
Woosuk Kwon committed
643

644

bnellnm's avatar
bnellnm committed
645
def get_nvcc_cuda_version() -> Version:
646
647
648
649
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
650
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
651
652
653
    nvcc_output = subprocess.check_output(
        [CUDA_HOME + "/bin/nvcc", "-V"], universal_newlines=True
    )
654
655
656
657
658
659
    output = nvcc_output.split()
    release_idx = output.index("release") + 1
    nvcc_cuda_version = parse(output[release_idx].split(",")[0])
    return nvcc_cuda_version


zhuwenwen's avatar
zhuwenwen committed
660
661
662
663
664
665
def get_sha(root: Union[str, Path]) -> str:
    try:
        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root).decode('ascii').strip()
    except Exception:
        return 'Unknown'

666

zhuwenwen's avatar
zhuwenwen committed
667
def get_version_add(sha: Optional[str] = None) -> str:
zhuwenwen's avatar
zhuwenwen committed
668
    command = "git config --global --add safe.directory "+pwd  
zhuwenwen's avatar
zhuwenwen committed
669
    subprocess.run(command, shell=True, capture_output=False, text=True) 
zhuwenwen's avatar
zhuwenwen committed
670
671
    vllm_root = os.path.dirname(os.path.abspath(__file__))
    add_version_path = os.path.join(os.path.join(vllm_root, "vllm"), "version.py")
zhuwenwen's avatar
zhuwenwen committed
672
    major, minor, _ = torch.__version__.split('.')
673
674
675
676
    if add_git_version:
        if sha != 'Unknown':
            if sha is None:
                sha = get_sha(vllm_root)
zhuwenwen's avatar
zhuwenwen committed
677
            version = 'das.opt1.alpha.' + sha[:7]
678
    else:
zhuwenwen's avatar
zhuwenwen committed
679
        version = 'das.opt1.alpha'
680

zhuwenwen's avatar
zhuwenwen committed
681
682
683
684
685
686
687

    # dtk version
    if os.getenv("ROCM_PATH"):
        rocm_path = os.getenv('ROCM_PATH', "")
        rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version")
        with open(rocm_version_path, 'r',encoding='utf-8') as file:
            lines = file.readlines()
zhuwenwen's avatar
zhuwenwen committed
688
        rocm_version=lines[0].replace(".", "")
zhuwenwen's avatar
zhuwenwen committed
689
        version += ".dtk" + rocm_version
zhuwenwen's avatar
zhuwenwen committed
690
    
691
    new_version_content = f"""
692
try:
693
694
695
    __version__ = "0.13.0"
    __version_tuple__ = (0, 13, 0)
    __hcu_version__ = f'0.13.0+{version}' 
zhuwenwen's avatar
zhuwenwen committed
696
    
zhuwenwen's avatar
zhuwenwen committed
697
    from vllm.version import __version__, __version_tuple__, __hcu_version__
698
except Exception as e:
zhuwenwen's avatar
zhuwenwen committed
699
700
    import warnings

zhuwenwen's avatar
zhuwenwen committed
701
    warnings.warn(f"Failed to read commit hash:\\n + str(e)",
702
703
                  RuntimeWarning,
                  stacklevel=2)
704
705
    __version__ = "dev"
    __version_tuple__ = (0, 0, __version__)
zhuwenwen's avatar
zhuwenwen committed
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
    
    
def _prev_minor_version_was(version_str):
    '''Check whether a given version matches the previous minor version.

    Return True if version_str matches the previous minor version.

    For example - return True if the current version if 0.7.4 and the
    supplied version_str is '0.6'.

    Used for --show-hidden-metrics-for-version.
    '''
    # Match anything if this is a dev tree
    if __version_tuple__[0:2] == (0, 0):
        return True

    # Note - this won't do the right thing when we release 1.0!
723
    # assert __version_tuple__[0] == 0
zhuwenwen's avatar
zhuwenwen committed
724
725
726
727
728
729
730
731
732
    assert isinstance(__version_tuple__[1], int)
    return version_str == f"{{__version_tuple__[0]}}.{{__version_tuple__[1] - 1}}"


def _prev_minor_version():
    '''For the purpose of testing, return a previous minor version number.'''
    # In dev tree, this will return "0.-1", but that will work fine"
    assert isinstance(__version_tuple__[1], int)
    return f"{{__version_tuple__[0]}}.{{__version_tuple__[1] - 1}}"
733
"""
zhuwenwen's avatar
zhuwenwen committed
734
    
zhuwenwen's avatar
zhuwenwen committed
735
    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
736
        file.write(new_version_content)
zhuwenwen's avatar
zhuwenwen committed
737
738
739
740
741
742
743
744
    file.close()
    
    
def get_version():
    get_version_add()
    version_file = 'vllm/version.py'
    with open(version_file, encoding='utf-8') as f:
        exec(compile(f.read(), version_file, 'exec'))
zhuwenwen's avatar
zhuwenwen committed
745
    return locals()['__hcu_version__']
zhuwenwen's avatar
zhuwenwen committed
746
747


748
749
750
751
752
753
754
755
def get_gaudi_sw_version():
    """
    Returns the driver version.
    """
    # Enable console printing for `hl-smi` check
    output = subprocess.run("hl-smi",
                            shell=True,
                            text=True,
756
                            capture_output=True,
757
758
759
760
761
762
763
                            env={"ENABLE_CONSOLE": "true"})
    if output.returncode == 0 and output.stdout:
        return output.stdout.split("\n")[2].replace(
            " ", "").split(":")[1][:-1].split("-")[0]
    return "0.0.0"  # when hl-smi is not available


764
def get_vllm_version() -> str:
765
766
767
    # 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"):
768
769
770
        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")
771

772
773
    # version = get_version(write_to="vllm/_version.py")
    # sep = "+" if "+" not in version else "."  # dev versions might contain +
774

775
    if _no_device():
776
        if envs.VLLM_TARGET_DEVICE == "empty":
777
            version += f"{sep}empty"
778
    elif _is_cuda():
779
        if envs.VLLM_USE_PRECOMPILED and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX:
780
            version += f"{sep}precompiled"
781
782
        else:
            cuda_version = str(get_nvcc_cuda_version())
783
            if cuda_version != envs.VLLM_MAIN_CUDA_VERSION:
784
785
786
787
                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}"
788
    elif _is_hip():
789
        # Get the Rocm Version
zhuwenwen's avatar
zhuwenwen committed
790
        # rocm_version = get_rocm_version() or torch.version.hip
791
        # if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION:
zhuwenwen's avatar
zhuwenwen committed
792
        #     version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
zhuwenwen's avatar
zhuwenwen committed
793
        version = get_version()
794
    elif _is_tpu():
795
        version += f"{sep}tpu"
796
    elif _is_cpu():
797
798
        if envs.VLLM_TARGET_DEVICE == "cpu":
            version += f"{sep}cpu"
799
    elif _is_xpu():
800
        version += f"{sep}xpu"
801
    else:
802
        raise RuntimeError("Unknown runtime environment")
803

804
805
806
    return version


807
def get_requirements() -> list[str]:
808
    """Get Python package dependencies from requirements.txt."""
809
    requirements_dir = ROOT_DIR / "requirements"
810

811
    def _read_requirements(filename: str) -> list[str]:
812
        with open(requirements_dir / filename) as f:
813
            requirements = f.read().strip().split("\n")
814
815
816
817
        resolved_requirements = []
        for line in requirements:
            if line.startswith("-r "):
                resolved_requirements += _read_requirements(line.split()[1])
818
819
820
821
822
            elif (
                not line.startswith("--")
                and not line.startswith("#")
                and line.strip() != ""
            ):
823
824
825
                resolved_requirements.append(line)
        return resolved_requirements

826
    if _no_device():
827
        requirements = _read_requirements("common.txt")
828
    elif _is_cuda():
829
        requirements = _read_requirements("cuda.txt")
830
        cuda_major, cuda_minor = torch.version.cuda.split(".")
831
832
        modified_requirements = []
        for req in requirements:
833
            if "vllm-flash-attn" in req and cuda_major != "12":
834
                # vllm-flash-attn is built only for CUDA 12.x.
835
836
837
                # Skip for other versions.
                continue
            modified_requirements.append(req)
838
        requirements = modified_requirements
839
    elif _is_hip():
840
        requirements = _read_requirements("rocm.txt")
841
    elif _is_tpu():
842
        requirements = _read_requirements("tpu.txt")
843
    elif _is_cpu():
844
        requirements = _read_requirements("cpu.txt")
845
    elif _is_xpu():
846
        requirements = _read_requirements("xpu.txt")
847
    else:
848
        raise ValueError("Unsupported platform, please use CUDA, ROCm, or CPU.")
849
850
851
    return requirements


bnellnm's avatar
bnellnm committed
852
853
ext_modules = []

854
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
855
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))
856
    ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
857
858
    # Optional since this doesn't get built (produce an .so file). This is just
    # copying the relevant .py files from the source repository.
859
    # ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
bnellnm's avatar
bnellnm committed
860

861
862
# if _is_hip():
#     ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
863

864
if _is_cuda():
865
    ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C"))
866
867
    if envs.VLLM_USE_PRECOMPILED or get_nvcc_cuda_version() >= Version("12.3"):
        # FA3 requires CUDA 12.3 or later
868
        ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
869
870
        # Optional since this doesn't get built (produce an .so file) when
        # not targeting a hopper system
871
        ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
872
        ext_modules.append(
873
874
            CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
        )
875

876
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
877
878
    ext_modules.append(CMakeExtension(name="vllm._C"))

zhuwenwen's avatar
zhuwenwen committed
879
880
881
882
883
884
885
886
887
if skip_vllm_build:
    package_data = {
        "vllm": [
            "py.typed",
            "model_executor/layers/fused_moe/configs/*.json",
            "model_executor/layers/quantization/utils/configs/*.json",
            "perf/*.py",
            "attention/backends/configs/*.json",
            "model_executor/layers/quantization/configs/awq/*.json",
zhuwenwen's avatar
zhuwenwen committed
888
889
            "_C.abi3.so",
            "_moe_C.abi3.so",
zhuwenwen's avatar
zhuwenwen committed
890
891
892
893
894
895
896
897
898
899
900
901
902
903
        ]
    }
else:
    package_data = {
        "vllm": [
            "py.typed",
            "model_executor/layers/fused_moe/configs/*.json",
            "model_executor/layers/quantization/utils/configs/*.json",
            "perf/*.py",
            "attention/backends/configs/*.json",
            "model_executor/layers/quantization/configs/awq/*.json",
        ]
    }

904

905
906
# If using precompiled, extract and patch package_data (in advance of setup)
if envs.VLLM_USE_PRECOMPILED:
907
    wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()
908
    patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(
909
910
        wheel_url, download_filename
    )
911
912
913
    for pkg, files in patch.items():
        package_data.setdefault(pkg, []).extend(files)

zhuwenwen's avatar
zhuwenwen committed
914
if _no_device() or skip_vllm_build:
915
916
    ext_modules = []

917
918
919
920
if not ext_modules:
    cmdclass = {}
else:
    cmdclass = {
921
922
923
        "build_ext": precompiled_build_ext
        if envs.VLLM_USE_PRECOMPILED
        else cmake_build_ext
924
925
    }

bnellnm's avatar
bnellnm committed
926
setup(
927
    # static metadata should rather go in pyproject.toml
928
    version=get_vllm_version(),
Woosuk Kwon's avatar
Woosuk Kwon committed
929
    ext_modules=ext_modules,
930
    install_requires=get_requirements(),
931
    extras_require={
932
        "bench": ["pandas", "matplotlib", "seaborn", "datasets"],
933
        "tensorizer": ["tensorizer==2.10.1"],
934
        "fastsafetensors": ["fastsafetensors >= 0.1.10"],
935
        "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
936
937
938
939
940
        "audio": [
            "librosa",
            "soundfile",
            "mistral_common[audio]",
        ],  # Required for audio processing
941
        "video": [],  # Kept for backwards compatibility
942
        "flashinfer": [],  # Kept for backwards compatibility
943
944
        # Optional deps for AMD FP4 quantization support
        "petit-kernel": ["petit-kernel"],
945
    },
946
    cmdclass=cmdclass,
Simon Mo's avatar
Simon Mo committed
947
    package_data=package_data,
Woosuk Kwon's avatar
Woosuk Kwon committed
948
)