setup.py

# SPDX-License-Identifier: Apache-2.0

import importlib.util
import logging
import os
import shutil
import subprocess
import sys
from pathlib import Path
from shutil import which

from packaging.version import Version
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_scm import get_version
from torch.utils.cpp_extension import SYCL_HOME


def load_module_from_path(module_name, path):
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


ROOT_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)

envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'tools',
                                                  'envs.py'))

VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE


def is_sccache_available() -> bool:
    return which("sccache") is not None


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


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


def is_url_available(url: str) -> bool:
    from urllib.request import urlopen

    status = None
    try:
        with urlopen(url) as f:
            status = f.status
    except Exception:
        return False
    return status == 200


def get_oneapi_version() -> Version:
    """Get the oneapi version from
    """
    assert SYCL_HOME is not None, "SYCL_HOME environment variable is not set."
    icpx_output = subprocess.check_output([SYCL_HOME + "/bin/icpx", "-v"],
                                          universal_newlines=True)
    print("=============== icpx version ===============")
    print(f"sycl home: {SYCL_HOME}")
    print(icpx_output)
    print("=============== icpx version ===============")
    # output = icpx_output.split()


def _build_custom_ops() -> bool:
    return True


class CMakeExtension(Extension):

    def __init__(self, name: str, cmake_lists_dir: str = '.', **kwa) -> None:
        super().__init__(name, sources=[], py_limited_api=True, **kwa)
        self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)


class cmake_build_ext(build_ext):
    # A dict of extension directories that have been configured.
    did_config: dict[str, bool] = {}

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
        # `num_jobs` is either the value of the MAX_JOBS environment variable
        # (if defined) or the number of CPUs available.
        num_jobs = envs.MAX_JOBS
        if num_jobs is not None:
            num_jobs = int(num_jobs)
            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
        else:
            try:
                # os.sched_getaffinity() isn't universally available, so fall
                #  back to os.cpu_count() if we get an error here.
                num_jobs = len(os.sched_getaffinity(0))
            except AttributeError:
                num_jobs = os.cpu_count()

        get_oneapi_version()

        return num_jobs

    #
    # 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 "Release"
        cfg = envs.CMAKE_BUILD_TYPE or default_cfg

        cmake_args = [
            '-DCMAKE_BUILD_TYPE={}'.format(cfg),
            '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
            '-DCMAKE_TOOLCHAIN_FILE=cmake/toolchain.cmake'
        ]

        verbose = envs.VERBOSE
        if verbose:
            cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
            ]
        elif is_ccache_available():
            cmake_args += [
                '-DCMAKE_C_COMPILER_LAUNCHER=ccache',
                '-DCMAKE_CXX_COMPILER_LAUNCHER=ccache',
            ]

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

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

        # Override the base directory for FetchContent downloads to $ROOT/.deps
        # This allows sharing dependencies between profiles,
        # and plays more nicely with sccache.
        # To override this, set the FETCHCONTENT_BASE_DIR environment variable.
        fc_base_dir = os.path.join(ROOT_DIR, ".deps")
        fc_base_dir = os.environ.get("FETCHCONTENT_BASE_DIR", fc_base_dir)
        cmake_args += ['-DFETCHCONTENT_BASE_DIR={}'.format(fc_base_dir)]

        #
        # Setup parallelism and build tool
        #
        num_jobs = self.compute_num_jobs()

        if is_ninja_available():
            build_tool = ['-G', 'Ninja']
            cmake_args += [
                '-DCMAKE_JOB_POOL_COMPILE:STRING=compile',
                '-DCMAKE_JOB_POOLS:STRING=compile={}'.format(num_jobs),
            ]
        else:
            # Default build tool to whatever cmake picks.
            build_tool = []
        my_env = os.environ.copy()
        icx_path = shutil.which('icx')
        icpx_path = shutil.which('icpx')
        build_option_gpu = {
            "CMAKE_C_COMPILER": f"{icx_path}",
            "CMAKE_CXX_COMPILER": f"{icpx_path}",
        }
        for key, value in build_option_gpu.items():
            if value is not None:
                cmake_args.append("-D{}={}".format(key, value))
        subprocess.check_call(
            ['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args],
            cwd=self.build_temp,
            env=my_env)

    def build_extensions(self) -> None:
        # Ensure that CMake is present and working
        try:
            subprocess.check_output(['cmake', '--version'])
        except OSError as e:
            raise RuntimeError('Cannot find CMake executable') from e

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

        targets = []

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

        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
            targets.append(target_name(ext.name))

        num_jobs = self.compute_num_jobs()

        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]

        subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)

        # Install the libraries
        for ext in self.extensions:
            # Install the extension into the proper location
            outdir = Path(self.get_ext_fullpath(ext.name)).parent.absolute()

            # Skip if the install directory is the same as the build directory
            if outdir == self.build_temp:
                continue

            # CMake appends the extension prefix to the install path,
            # and outdir already contains that prefix, so we need to remove it.
            # We assume only the final component of extension prefix is added by
            # CMake, this is currently true for current extensions but may not
            # always be the case.
            prefix = outdir
            if '.' in ext.name:
                prefix = prefix.parent
            # prefix here should actually be the same for all components
            install_args = [
                "cmake", "--install", ".", "--prefix", prefix, "--component",
                target_name(ext.name)
            ]
            subprocess.check_call(install_args, cwd=self.build_temp)

        # Install additional shared libraries (intermediate build artifacts)
        # These are compiled as separate libraries but need to be packaged in
        # the wheel
        if self.extensions:
            # Use the same prefix as the extensions
            first_ext = self.extensions[0]
            outdir = Path(self.get_ext_fullpath(
                first_ext.name)).parent.absolute()
            prefix = outdir.parent if '.' in first_ext.name else outdir

            for lib_name, file_path in additional_libraries.items():
                install_args = [
                    "cmake",
                    "--install",
                    ".",
                    "--prefix",
                    prefix,
                    "--component",
                    lib_name,
                    "--verbose",
                ]
                try:
                    subprocess.check_call(install_args,
                                          cwd=self.build_temp + file_path)
                except subprocess.CalledProcessError as e:
                    logger.warning("Failed to install library %s: %s",
                                   lib_name, e)
                    # Continue with other libraries even if one fails

    def run(self):
        self.build_temp = "build/temp"
        # First, run the standard build_ext command to compile the extensions
        super().run()

        import glob
        files = glob.glob(
            os.path.join(self.build_lib, "vllm_xpu_kernels", "lib*.so"))
        # if is editable install, also copy to local inplace directory
        if self.inplace:
            for file in files:
                inplace_dst_file = os.path.join(
                    os.path.dirname(__file__),
                    "vllm_xpu_kernels",
                    file.split("vllm_xpu_kernels/")[-1],
                )
                print(f"Copying {file} to {inplace_dst_file}")
                self.copy_file(file, inplace_dst_file)


ext_modules = []

# List of additional shared libraries to install (intermediate build artifacts)
additional_libraries = {
    "attn_kernels_xe_2": "/csrc/xpu/attn/xe_2",
    "gdn_attn_kernels_xe_2": "/csrc/xpu/gdn_attn/xe_2",
    "grouped_gemm_xe_default": "/csrc/xpu/grouped_gemm/xe_default",
    "grouped_gemm_xe_2": "/csrc/xpu/grouped_gemm/xe_2",
}

if _build_custom_ops():
    ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._C"))
    ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._vllm_fa2_C"))
    ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._moe_C"))
    ext_modules.append(CMakeExtension(name="vllm_xpu_kernels._xpu_C"))
    ext_modules.append(
        CMakeExtension(name="vllm_xpu_kernels.xpumem_allocator"))

if ext_modules:
    cmdclass = {"build_ext": cmake_build_ext}

package_data = {
    "vllm-xpu-kernels": [
        "py.typed",
    ]
}


def get_vllm_version() -> str:
    # Allow overriding the version.
    if env_version := os.getenv("VLLM_VERSION_OVERRIDE"):
        print(f"Overriding VLLM version with {env_version}")
        os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version
        return get_version(write_to="_version.py")

    version = get_version(write_to="_version.py")

    return version


setup(
    version=get_vllm_version(),
    ext_modules=ext_modules,
    cmdclass=cmdclass,
    package_data=package_data,
)