ref(upstream): FULL TREE — Deep-Spark xllm (1470) + ds_vllm csrc/models (703)
Replaces cherry-picked upstream_ref with complete source trees. xllm/ — Iluvatar official C++ inference engine (15MB, 1470 files) Complete: kernels → layers → models → runtime → scheduler → api Excluded: .git, binary images, third_party submodule checkouts ds_vllm/ — Iluvatar official vllm fork (8MB, 703 files) Included: csrc/ (ALL CUDA kernels), fused_moe/, qwen3_5 model, _custom_ops Excluded: tests, benchmarks, docs, examples (not needed for reference) Critical call chains now fully traceable: MoE: moe_topk_softmax_kernels.cuh → ixformer.h → fused_moe.cpp → layer GDN: qwen3_gated_delta_net_base.cpp → qwen3_5_gated_delta_net.cpp Attention: ixformer.h → xllm_paged_attention → attention.cpp
This commit is contained in:
EX Engine
670
upstream_ref/xllm/setup.py
Normal file
670
upstream_ref/xllm/setup.py
Normal file
@@ -0,0 +1,670 @@
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import argparse
|
||||
from typing import Any, Optional
|
||||
|
||||
from distutils.core import Command
|
||||
from setuptools import Extension, find_namespace_packages, setup
|
||||
from setuptools.command.build_ext import build_ext
|
||||
|
||||
try:
|
||||
from setuptools.command.bdist_wheel import bdist_wheel
|
||||
except ModuleNotFoundError:
|
||||
from wheel.bdist_wheel import bdist_wheel
|
||||
|
||||
from scripts.build_support.env import (
|
||||
get_cxx_abi,
|
||||
set_cuda_envs,
|
||||
set_ilu_envs,
|
||||
set_mlu_envs,
|
||||
set_musa_envs,
|
||||
set_npu_envs,
|
||||
)
|
||||
from scripts.build_support.utils import (
|
||||
check_and_install_pre_commit,
|
||||
get_base_dir,
|
||||
get_cmake_dir,
|
||||
get_cpu_arch,
|
||||
get_device_type,
|
||||
get_python_version,
|
||||
get_torch_version,
|
||||
get_version,
|
||||
pre_build,
|
||||
read_readme,
|
||||
)
|
||||
|
||||
BUILD_TEST_FILE: bool = True
|
||||
BUILD_EXPORT: bool = True
|
||||
|
||||
|
||||
def _maybe_compile_tilelang_kernels(device: str) -> None:
|
||||
if device != "npu":
|
||||
return
|
||||
|
||||
output_root = os.path.join(get_cmake_dir(), "xllm", "compiler", "tilelang")
|
||||
os.makedirs(output_root, exist_ok=True)
|
||||
|
||||
env = os.environ.copy()
|
||||
base_dir = get_base_dir()
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
os.path.join(base_dir, "xllm", "compiler", "tilelang_launcher.py"),
|
||||
"compile-kernels",
|
||||
"--target",
|
||||
"ascend",
|
||||
"--output-root",
|
||||
output_root,
|
||||
]
|
||||
print("[INFO] compiling TileLang kernels via source-tree launcher")
|
||||
subprocess.check_call(cmd, cwd=base_dir, env=env)
|
||||
|
||||
class CMakeExtension(Extension):
|
||||
def __init__(self, name: str, path: str, sourcedir: str = "") -> None:
|
||||
super().__init__(name, sources=[])
|
||||
self.sourcedir = os.path.realpath(os.path.abspath(sourcedir))
|
||||
self.path = path
|
||||
|
||||
class ExtBuild(build_ext):
|
||||
user_options = build_ext.user_options + [
|
||||
("base-dir=", None, "base directory of xLLM project"),
|
||||
("device=", None, "target device type (npu or mlu or cuda or ilu or musa)"),
|
||||
("arch=", None, "target arch type (x86 or arm)"),
|
||||
("generate-so=", None, "generate so or binary"),
|
||||
]
|
||||
|
||||
def initialize_options(self) -> None:
|
||||
build_ext.initialize_options(self)
|
||||
self.base_dir = get_base_dir()
|
||||
self.device: Optional[str] = None
|
||||
self.arch: Optional[str] = None
|
||||
self.generate_so: bool = False
|
||||
|
||||
def finalize_options(self) -> None:
|
||||
build_ext.finalize_options(self)
|
||||
|
||||
def run(self) -> None:
|
||||
# check if cmake is installed
|
||||
try:
|
||||
out: bytes = subprocess.check_output(["cmake", "--version"])
|
||||
except OSError:
|
||||
raise RuntimeError(
|
||||
"CMake must be installed to build the following extensions: "
|
||||
+ ", ".join(e.name for e in self.extensions)
|
||||
)
|
||||
exit(1)
|
||||
|
||||
match = re.search(
|
||||
r"version\s*(?P<major>\d+)\.(?P<minor>\d+)([\d.]+)?", out.decode()
|
||||
)
|
||||
if match is None:
|
||||
raise RuntimeError(f"Failed to parse CMake version from: {out!r}")
|
||||
cmake_major, cmake_minor = int(match.group("major")), int(match.group("minor"))
|
||||
if (cmake_major, cmake_minor) < (3, 18):
|
||||
raise RuntimeError("CMake >= 3.18.0 is required")
|
||||
|
||||
try:
|
||||
# build extensions
|
||||
for ext in self.extensions:
|
||||
self.build_extension(ext)
|
||||
except Exception as e:
|
||||
print("ERROR: Build failed.")
|
||||
print(f"Details: {e}")
|
||||
exit(1)
|
||||
|
||||
def build_extension(self, ext: CMakeExtension) -> None:
|
||||
ninja_dir = shutil.which("ninja")
|
||||
# the output dir for the extension
|
||||
extdir: str = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
|
||||
|
||||
# create build directory
|
||||
os.makedirs(self.build_temp, exist_ok=True)
|
||||
|
||||
# Using this requires trailing slash for auto-detection & inclusion of
|
||||
# auxiliary "native" libs
|
||||
|
||||
debug: int = int(os.environ.get("DEBUG", 0)) if self.debug is None else int(self.debug)
|
||||
build_type: str = "Debug" if debug else "Release"
|
||||
|
||||
default_jobs = os.cpu_count() or 1
|
||||
max_jobs: str = os.getenv("MAX_JOBS", str(default_jobs))
|
||||
max_jobs_int: int = int(max_jobs)
|
||||
|
||||
# Limit archive (ar/ranlib) concurrency to avoid file locking conflicts.
|
||||
# The ar tool requires exclusive access to archive files (.a files) when
|
||||
# creating or updating static libraries. When multiple ar processes attempt
|
||||
# to modify the same archive file simultaneously, they compete for file locks,
|
||||
# which can cause deadlocks and hang the build process.
|
||||
archive_jobs: int = min(8, max(1, max_jobs_int // 4))
|
||||
cmake_args: list[str] = [
|
||||
"-G",
|
||||
"Ninja",
|
||||
f"-DCMAKE_MAKE_PROGRAM={ninja_dir}",
|
||||
f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
|
||||
f"-DCMAKE_RUNTIME_OUTPUT_DIRECTORY={extdir}",
|
||||
"-DUSE_CCACHE=ON",
|
||||
f"-DPython_EXECUTABLE:FILEPATH={sys.executable}",
|
||||
f"-DCMAKE_BUILD_TYPE={build_type}",
|
||||
f"-DBUILD_SHARED_LIBS=OFF",
|
||||
f"-DDEVICE_TYPE=USE_{self.device.upper()}",
|
||||
f"-DDEVICE_ARCH={self.arch.upper()}",
|
||||
f"-DXLLM_ATB_LAYERS_SOURCE_DIR={os.path.join(self.base_dir, 'third_party', 'xllm_atb_layers')}",
|
||||
f"-DCMAKE_JOB_POOLS=archive={archive_jobs}",
|
||||
]
|
||||
|
||||
if self.device is None:
|
||||
raise ValueError("Please set --device to npu or mlu or cuda or ilu or musa.")
|
||||
if self.arch is None:
|
||||
raise ValueError("Please set --arch to x86 or arm.")
|
||||
|
||||
if self.device == "npu":
|
||||
cmake_args += ["-DUSE_NPU=ON"]
|
||||
set_npu_envs()
|
||||
_maybe_compile_tilelang_kernels(self.device)
|
||||
elif self.device == "mlu":
|
||||
cmake_args += ["-DUSE_MLU=ON"]
|
||||
set_mlu_envs()
|
||||
elif self.device == "cuda":
|
||||
torch_cuda_architectures = os.getenv("TORCH_CUDA_ARCH_LIST")
|
||||
if not torch_cuda_architectures:
|
||||
raise ValueError("Please set TORCH_CUDA_ARCH_LIST environment variable, e.g. export TORCH_CUDA_ARCH_LIST=\"8.0 8.9 9.0 10.0 12.0\"")
|
||||
cmake_args += ["-DUSE_CUDA=ON",
|
||||
f"-DTORCH_CUDA_ARCH_LIST={torch_cuda_architectures}"]
|
||||
set_cuda_envs()
|
||||
elif self.device == "ilu":
|
||||
cmake_args += ["-DUSE_ILU=ON"]
|
||||
set_ilu_envs()
|
||||
elif self.device == "musa":
|
||||
cmake_args += ["-DUSE_MUSA=ON"]
|
||||
set_musa_envs()
|
||||
global BUILD_TEST_FILE
|
||||
BUILD_TEST_FILE = False
|
||||
else:
|
||||
raise ValueError("Please set --device to npu or mlu or cuda or ilu or musa.")
|
||||
|
||||
product: str = "xllm"
|
||||
if self.generate_so:
|
||||
product = "libxllm.so"
|
||||
cmake_args += ["-DGENERATE_SO=ON"]
|
||||
else:
|
||||
cmake_args += ["-DGENERATE_SO=OFF"]
|
||||
|
||||
# Adding CMake arguments set as environment variable
|
||||
# (needed e.g. to build for ARM OSx on conda-forge)
|
||||
if "CMAKE_ARGS" in os.environ:
|
||||
cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
|
||||
|
||||
# check if torch binary is built with cxx11 abi
|
||||
if get_cxx_abi():
|
||||
cmake_args += ["-DUSE_CXX11_ABI=ON", "-D_GLIBCXX_USE_CXX11_ABI=1"]
|
||||
else:
|
||||
cmake_args += ["-DUSE_CXX11_ABI=OFF", "-D_GLIBCXX_USE_CXX11_ABI=0"]
|
||||
|
||||
build_args = ["--config", build_type]
|
||||
build_args += ["-j" + max_jobs]
|
||||
|
||||
env: dict[str, str] = os.environ.copy()
|
||||
env["VCPKG_MAX_CONCURRENCY"] = str(max_jobs)
|
||||
print("CMake Args: ", cmake_args)
|
||||
print("Env: ", env)
|
||||
|
||||
self.build_cmake_targets(ext, cmake_args, build_args, env, extdir, product)
|
||||
|
||||
def build_cmake_targets(
|
||||
self,
|
||||
ext: CMakeExtension,
|
||||
cmake_args: list[str],
|
||||
build_args: list[str],
|
||||
env: dict[str, str],
|
||||
extdir: str,
|
||||
product: str,
|
||||
) -> None:
|
||||
"""Build CMake targets"""
|
||||
cmake_dir = get_cmake_dir()
|
||||
subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)
|
||||
|
||||
base_build_args = build_args
|
||||
# add build target to speed up the build process
|
||||
build_args += ["--target", ext.name, "xllm"]
|
||||
subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)
|
||||
|
||||
os.makedirs(os.path.join(os.path.dirname(cmake_dir), "xllm/core/server/"), exist_ok=True)
|
||||
shutil.copy(
|
||||
os.path.join(extdir, product),
|
||||
os.path.join(os.path.dirname(cmake_dir), "xllm/core/server/"),
|
||||
)
|
||||
|
||||
if BUILD_EXPORT:
|
||||
# build export module
|
||||
build_args = base_build_args + ["--target export_module"]
|
||||
subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)
|
||||
|
||||
if BUILD_TEST_FILE:
|
||||
# build tests target
|
||||
build_args = base_build_args + ["--target all_tests"]
|
||||
subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)
|
||||
|
||||
class ExtBuildSingleTest(ExtBuild):
|
||||
"""Inherit ExtBuild, used to build and run a single test"""
|
||||
user_options = ExtBuild.user_options + [
|
||||
("test-name=", None, "name of the test target to build and run"),
|
||||
]
|
||||
|
||||
def initialize_options(self) -> None:
|
||||
ExtBuild.initialize_options(self)
|
||||
self.test_name: Optional[str] = None
|
||||
|
||||
def finalize_options(self) -> None:
|
||||
ExtBuild.finalize_options(self)
|
||||
if not self.test_name:
|
||||
raise ValueError("--test-name is required for ExtBuildSingleTest")
|
||||
|
||||
def build_cmake_targets(
|
||||
self,
|
||||
ext: CMakeExtension,
|
||||
cmake_args: list[str],
|
||||
build_args: list[str],
|
||||
env: dict[str, str],
|
||||
extdir: str,
|
||||
product: str,
|
||||
) -> None:
|
||||
"""Override method: only build the specified test target and run"""
|
||||
cmake_dir = get_cmake_dir()
|
||||
subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)
|
||||
|
||||
base_build_args = build_args
|
||||
# Only build the specified test target
|
||||
build_args += ["--target", self.test_name]
|
||||
subprocess.check_call(["cmake", "--build", ".", "--verbose"] + build_args, cwd=cmake_dir)
|
||||
|
||||
# Find test executable
|
||||
# CMake usually places executables in CMAKE_RUNTIME_OUTPUT_DIRECTORY or build directory
|
||||
test_executable: Optional[str] = None
|
||||
possible_paths: list[str] = [
|
||||
os.path.join(cmake_dir, self.test_name),
|
||||
os.path.join(extdir, self.test_name),
|
||||
os.path.join(cmake_dir, "xllm", "core", self.test_name),
|
||||
]
|
||||
|
||||
# Check possible paths first
|
||||
for path in possible_paths:
|
||||
if os.path.exists(path) and os.access(path, os.X_OK):
|
||||
test_executable = path
|
||||
break
|
||||
|
||||
# If not found, try recursive search in build directory
|
||||
if not test_executable:
|
||||
for root, dirs, files in os.walk(cmake_dir):
|
||||
if self.test_name in files:
|
||||
candidate = os.path.join(root, self.test_name)
|
||||
if os.access(candidate, os.X_OK):
|
||||
test_executable = candidate
|
||||
break
|
||||
|
||||
if not test_executable:
|
||||
# If not found, try using ctest to run
|
||||
print(f"⚠️ Warning: Could not find test executable {self.test_name}, trying ctest...")
|
||||
try:
|
||||
subprocess.check_call(
|
||||
["ctest", "-R", self.test_name, "--verbose"],
|
||||
cwd=cmake_dir,
|
||||
env=env
|
||||
)
|
||||
print(f"✅ Test {self.test_name} passed!")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"❌ Failed to run test {self.test_name}")
|
||||
raise
|
||||
else:
|
||||
# Run test executable directly
|
||||
print(f"🚀 Running test: {test_executable}")
|
||||
try:
|
||||
subprocess.check_call([test_executable], cwd=os.path.dirname(test_executable), env=env)
|
||||
print(f"✅ Test {self.test_name} passed!")
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"❌ Test {self.test_name} failed with exit code {e.returncode}")
|
||||
raise
|
||||
|
||||
class BuildDistWheel(bdist_wheel):
|
||||
user_options = bdist_wheel.user_options + [
|
||||
("device=", None, "target device type (npu or mlu or cuda or ilu or musa)"),
|
||||
("arch=", None, "target arch type (x86 or arm)"),
|
||||
]
|
||||
|
||||
def initialize_options(self) -> None:
|
||||
super().initialize_options()
|
||||
self.device: Optional[str] = None
|
||||
self.arch: Optional[str] = None
|
||||
# Cache the original dist name early so finalize_options is idempotent
|
||||
# and so name changes are visible to egg_info/metadata generation.
|
||||
self._base_dist_name = self.distribution.metadata.name
|
||||
|
||||
def finalize_options(self) -> None:
|
||||
# IMPORTANT: mutate distribution name BEFORE super().finalize_options().
|
||||
# bdist_wheel finalization may finalize/cache egg_info metadata; if we
|
||||
# change the name afterwards, the wheel filename and METADATA can diverge
|
||||
# (pip will reject the wheel as "inconsistent name").
|
||||
name = self._base_dist_name
|
||||
|
||||
# generate distribution name suffix
|
||||
if self.device:
|
||||
name += f"_{self.device}"
|
||||
|
||||
torch_version = get_torch_version(self.device)
|
||||
if torch_version:
|
||||
name += f"_torch{torch_version}"
|
||||
|
||||
if get_cxx_abi():
|
||||
name += "_cxx11_abi"
|
||||
else:
|
||||
name += "_no_cxx11_abi"
|
||||
|
||||
self.distribution.metadata.name = name
|
||||
super().finalize_options()
|
||||
|
||||
def run(self) -> None:
|
||||
build_ext_cmd = self.get_finalized_command('build_ext')
|
||||
build_ext_cmd.device = self.device
|
||||
build_ext_cmd.arch = self.arch
|
||||
|
||||
print("🔨 build project...")
|
||||
self.run_command('build')
|
||||
|
||||
print("🧪 testing UT...")
|
||||
self.run_command('test')
|
||||
|
||||
if self.arch == 'arm':
|
||||
ext_path = get_base_dir() + f"/build/lib.linux-aarch64-cpython-{get_python_version()}/"
|
||||
else:
|
||||
ext_path = get_base_dir() + f"/build/lib.linux-x86_64-cpython-{get_python_version()}/"
|
||||
if len(ext_path) == 0:
|
||||
print("❌ Build wheel failed, not found path.")
|
||||
exit(1)
|
||||
tmp_path = os.path.join(ext_path, 'xllm')
|
||||
for root, dirs, files in os.walk(tmp_path):
|
||||
for item in files:
|
||||
path = os.path.join(root, item)
|
||||
if '_test' in item and os.path.isfile(path):
|
||||
os.remove(path)
|
||||
global BUILD_TEST_FILE
|
||||
BUILD_TEST_FILE = False
|
||||
|
||||
self.skip_build = True
|
||||
super().run()
|
||||
|
||||
class TestUT(Command):
|
||||
description = "Run all testing binary."
|
||||
user_options = []
|
||||
|
||||
# Whitelist: tests that must run sequentially (not in parallel with others)
|
||||
# Add test names here if they use fork() or have device initialization conflicts
|
||||
# Note: Use test case name patterns (from gtest), not executable names
|
||||
SEQUENTIAL_TESTS = [
|
||||
'ReduceScatterMultiDeviceTest',
|
||||
'DeepEPMultiDeviceTest',
|
||||
'AttentionMultiDeviceTest',
|
||||
'FusedMoEAll2AllMultiDeviceTest',
|
||||
]
|
||||
|
||||
def initialize_options(self) -> None:
|
||||
pass
|
||||
|
||||
def finalize_options(self) -> None:
|
||||
pass
|
||||
|
||||
def run_ctest(self, cmake_dir: str) -> int:
|
||||
def run_subprocess_with_streaming(
|
||||
cmd: list[str],
|
||||
error_message: str,
|
||||
warn_if_no_tests: bool = False,
|
||||
) -> None:
|
||||
"""Helper function to run subprocess and stream output"""
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
cwd=cmake_dir,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
|
||||
if process.stdout is None:
|
||||
raise RuntimeError("Failed to capture subprocess stdout for streaming.")
|
||||
|
||||
output_lines: list[str] = []
|
||||
for line in iter(process.stdout.readline, ''):
|
||||
print(line, end='')
|
||||
output_lines.append(line)
|
||||
|
||||
return_code: int = process.wait()
|
||||
|
||||
# Warn if no tests were found, but don't fail (some backends may not compile certain tests)
|
||||
if warn_if_no_tests and return_code == 0:
|
||||
output_text: str = ''.join(output_lines)
|
||||
if 'No tests were found' in output_text:
|
||||
print(f"No tests matched the pattern (this is OK for some backends).")
|
||||
return
|
||||
|
||||
if return_code != 0:
|
||||
print(error_message)
|
||||
exit(1)
|
||||
|
||||
try:
|
||||
# Step 1: Run all tests EXCEPT sequential ones in parallel
|
||||
if self.SEQUENTIAL_TESTS:
|
||||
exclude_pattern = '|'.join(self.SEQUENTIAL_TESTS)
|
||||
print("=" * 80)
|
||||
print(f"Running tests in parallel (excluding: {', '.join(self.SEQUENTIAL_TESTS)})...")
|
||||
print("=" * 80)
|
||||
run_subprocess_with_streaming(
|
||||
['ctest', '--parallel', '8', '--repeat', 'until-pass:5', '-E', exclude_pattern],
|
||||
"Parallel tests failed."
|
||||
)
|
||||
else:
|
||||
print("=" * 80)
|
||||
print("Running all tests in parallel...")
|
||||
print("=" * 80)
|
||||
run_subprocess_with_streaming(
|
||||
['ctest', '--parallel', '8', '--repeat', 'until-pass:5'],
|
||||
"Parallel tests failed."
|
||||
)
|
||||
|
||||
# Step 2: Run sequential tests one by one
|
||||
for idx, test_name in enumerate(self.SEQUENTIAL_TESTS, start=2):
|
||||
print("\n" + "=" * 80)
|
||||
print(f"Step {idx}: Running {test_name} sequentially...")
|
||||
print("=" * 80)
|
||||
# Use pattern matching to include all test cases under the test class
|
||||
# e.g., ReduceScatterMultiDeviceTest matches ReduceScatterMultiDeviceTest.BasicTest, etc.
|
||||
run_subprocess_with_streaming(
|
||||
['ctest', '--repeat', 'until-pass:5', '-R', test_name],
|
||||
f"Sequential test {test_name} failed.",
|
||||
warn_if_no_tests=True
|
||||
)
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("All tests passed!")
|
||||
print("=" * 80)
|
||||
return 0
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(e.stderr)
|
||||
exit(1)
|
||||
|
||||
def run(self) -> None:
|
||||
self.run_ctest(get_cmake_dir())
|
||||
|
||||
class SingleTest(Command):
|
||||
"""Command to build and run a single test"""
|
||||
description = "Build and run a single test target."
|
||||
# test_name should match a CMake/CTest target name, for example:
|
||||
# python setup.py test --test-name common_test
|
||||
user_options = [
|
||||
("test-name=", None, "name of the test target to build and run (e.g. platform_vmm_test)"),
|
||||
("device=", None, "target device type (npu or mlu or cuda or ilu)"),
|
||||
("arch=", None, "target arch type (x86 or arm)"),
|
||||
("generate-so=", None, "generate so or binary"),
|
||||
]
|
||||
|
||||
def initialize_options(self) -> None:
|
||||
self.test_name: Optional[str] = None
|
||||
self.device: Optional[str] = None
|
||||
self.arch: Optional[str] = None
|
||||
self.generate_so: bool = False
|
||||
|
||||
def finalize_options(self) -> None:
|
||||
if not self.test_name:
|
||||
raise ValueError("--test-name is required for single_test command")
|
||||
|
||||
def run(self) -> None:
|
||||
# Create ExtBuildSingleTest instance and set parameters
|
||||
build_ext = ExtBuildSingleTest(self.distribution)
|
||||
build_ext.initialize_options()
|
||||
build_ext.test_name = self.test_name
|
||||
build_ext.device = self.device
|
||||
build_ext.arch = self.arch
|
||||
build_ext.generate_so = self.generate_so
|
||||
build_ext.finalize_options()
|
||||
|
||||
# Ensure extension modules are set
|
||||
if not hasattr(build_ext, 'extensions') or not build_ext.extensions:
|
||||
build_ext.extensions = self.distribution.ext_modules
|
||||
|
||||
# Run build
|
||||
build_ext.run()
|
||||
|
||||
def parse_arguments() -> dict[str, Any]:
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Setup helper for building xllm',
|
||||
epilog='Example: python setup.py build',
|
||||
usage='%(prog)s [COMMAND] [OPTIONS]'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'setup_args',
|
||||
nargs='*',
|
||||
metavar='argparse.REMAINDER',
|
||||
help='setup command (build, test, bdist_wheel, etc.)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--device',
|
||||
type=str.lower,
|
||||
choices=['auto', 'npu', 'mlu', 'cuda', 'ilu', 'musa'],
|
||||
default='auto',
|
||||
help='Device type: npu, mlu, ilu, cuda or musa (case-insensitive)'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--generate-so',
|
||||
type=str.lower,
|
||||
choices=['true', 'false', '1', '0', 'yes', 'no', 'y', 'n', 'on', 'off'],
|
||||
default='false',
|
||||
help='Whether to generate so or binary'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--test-name',
|
||||
type=str,
|
||||
default=None,
|
||||
help='Name of the test target to build and run; when omitted, all tests run'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
sys.argv = [sys.argv[0]] + args.setup_args
|
||||
|
||||
generate_so = args.generate_so.lower() in ('true', '1', 'yes', 'y', 'on')
|
||||
|
||||
return {
|
||||
'device': args.device,
|
||||
'generate_so': generate_so,
|
||||
'test_name': args.test_name,
|
||||
}
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = parse_arguments()
|
||||
|
||||
arch = get_cpu_arch()
|
||||
device = config['device']
|
||||
if device == 'auto':
|
||||
device = get_device_type()
|
||||
print(f"🚀 Build xllm with CPU arch: {arch} and target device: {device}")
|
||||
|
||||
pre_build()
|
||||
|
||||
generate_so = config['generate_so']
|
||||
test_name = config.get('test_name')
|
||||
|
||||
if "SKIP_TEST" in os.environ:
|
||||
BUILD_TEST_FILE = False
|
||||
if "SKIP_EXPORT" in os.environ:
|
||||
BUILD_EXPORT = False
|
||||
|
||||
version = get_version()
|
||||
|
||||
# check and install git pre-commit
|
||||
check_and_install_pre_commit()
|
||||
|
||||
test_cmd = SingleTest if test_name else TestUT
|
||||
options = {
|
||||
'build_ext': {
|
||||
'device': device,
|
||||
'arch': arch,
|
||||
'generate_so': generate_so
|
||||
},
|
||||
'bdist_wheel': {
|
||||
'device': device,
|
||||
'arch': arch,
|
||||
}
|
||||
}
|
||||
if test_name:
|
||||
options['test'] = {
|
||||
'device': device,
|
||||
'arch': arch,
|
||||
'generate_so': generate_so,
|
||||
'test_name': test_name,
|
||||
}
|
||||
|
||||
setup(
|
||||
name="xllm",
|
||||
version=version,
|
||||
license="Apache 2.0",
|
||||
author="xLLM Team",
|
||||
author_email="infer@xllm.ai",
|
||||
description="A high-performance inference system for large language models.",
|
||||
long_description=read_readme(),
|
||||
long_description_content_type="text/markdown",
|
||||
url="https://github.com/jd-opensource/xllm",
|
||||
project_urls={
|
||||
"Homepage": "https://xllm.readthedocs.io/zh-cn/latest/",
|
||||
"Documentation": "https://xllm.readthedocs.io/zh-cn/latest/",
|
||||
},
|
||||
classifiers=[
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Education",
|
||||
"Programming Language :: C++",
|
||||
"Programming Language :: Python :: 3 :: Only",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Operating System :: POSIX",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Topic :: Scientific/Engineering",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
],
|
||||
ext_modules=[CMakeExtension("xllm", "xllm/")],
|
||||
cmdclass={"build_ext": ExtBuild,
|
||||
"test": test_cmd,
|
||||
'bdist_wheel': BuildDistWheel},
|
||||
options=options,
|
||||
packages=find_namespace_packages(include=["scripts.build_support"]),
|
||||
zip_safe=False,
|
||||
py_modules=["xllm/launch_xllm", "xllm/__init__",
|
||||
"xllm/pybind/llm", "xllm/pybind/vlm",
|
||||
"xllm/pybind/embedding", "xllm/pybind/util",
|
||||
"xllm/pybind/args", "xllm/pybind/params",
|
||||
"xllm/pybind/errors", "xllm/pybind/mm_utils"],
|
||||
python_requires=">=3.10",
|
||||
)
|
||||
Reference in New Issue
Block a user