mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Derived from vLLM; modified by DeepSeek.
|
|
|
|
import os
|
|
|
|
import setuptools
|
|
import torch
|
|
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
|
|
|
|
from deep_cluster.hardware import accelerator_kind
|
|
|
|
IS_NPU = accelerator_kind() == "ascend"
|
|
|
|
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Extension configuration shared by the CUDA and NPU backends.
|
|
COMMON_COMPILE_FLAGS = ["-O2", "-std=c++20", "-Wno-deprecated-declarations"]
|
|
ABI_FLAG = "-D_GLIBCXX_USE_CXX11_ABI=1"
|
|
assert torch._C._GLIBCXX_USE_CXX11_ABI, "must use CXX11 ABI"
|
|
|
|
CXX_FLAGS = ["-g1", *COMMON_COMPILE_FLAGS, ABI_FLAG]
|
|
COMMON_SOURCES = [
|
|
"csrc/pybind.cpp",
|
|
"csrc/engine/pybind.cpp",
|
|
"csrc/engine/pybind_sequence.cpp",
|
|
"csrc/engine/pybind_scheduler.cpp",
|
|
"csrc/engine/pybind_output.cpp",
|
|
"csrc/engine/zos.cpp",
|
|
"csrc/engine/input_metadata/block_load_plan.cpp",
|
|
"csrc/engine/input_metadata/prepare_tensor_input_base.cpp",
|
|
"csrc/engine/input_metadata/prepare_tensor_input_v4.cpp",
|
|
"csrc/engine/prepare_sample.cpp",
|
|
"csrc/graph_timer/pybind.cpp",
|
|
]
|
|
COMMON_LIBRARIES = ["pthread"]
|
|
INCLUDE_DIRS = [os.path.abspath("csrc")]
|
|
|
|
if IS_NPU:
|
|
from torch_npu.utils.cpp_extension import NpuExtension # pyright: ignore[reportMissingImports]
|
|
|
|
deepinfer_extension = NpuExtension(
|
|
name="deepinfer._C",
|
|
sources=[
|
|
*COMMON_SOURCES,
|
|
"csrc/asc_utils.cpp",
|
|
"csrc/graph_timer/graph_timer_asc.cpp",
|
|
],
|
|
extra_compile_args={
|
|
"cxx": [*CXX_FLAGS, "-fopenmp", "-fvisibility=hidden", "-DWITH_NPU"],
|
|
},
|
|
libraries=COMMON_LIBRARIES,
|
|
include_dirs=INCLUDE_DIRS,
|
|
)
|
|
else:
|
|
conda_prefix = os.environ.get("CONDA_PREFIX") or os.environ.get("VIRTUAL_ENV")
|
|
if not conda_prefix:
|
|
raise RuntimeError("Neither CONDA_PREFIX nor VIRTUAL_ENV is set; a conda env or virtualenv is required to build.")
|
|
|
|
deepinfer_extension = CUDAExtension(
|
|
name="deepinfer._C",
|
|
sources=[
|
|
*COMMON_SOURCES,
|
|
"csrc/cuda_utils.cpp",
|
|
"csrc/graph_timer/graph_timer_cuda.cpp",
|
|
],
|
|
extra_compile_args={
|
|
"cxx": [*CXX_FLAGS, "-fvisibility=hidden"],
|
|
},
|
|
libraries=[*COMMON_LIBRARIES, "cuda"],
|
|
library_dirs=[os.path.join(CUDA_HOME, "lib64", "stubs")], # for libcuda.so
|
|
include_dirs=INCLUDE_DIRS,
|
|
)
|
|
|
|
|
|
def get_path(*filepath: str) -> str:
|
|
return os.path.join(ROOT_DIR, *filepath)
|
|
|
|
|
|
def get_requirements() -> list[str]:
|
|
"""Get Python package dependencies from requirements.txt (comments and blank lines stripped)."""
|
|
with open(get_path("requirements.txt")) as f:
|
|
lines = (line.split("#", 1)[0].strip() for line in f)
|
|
return [line for line in lines if line]
|
|
|
|
|
|
setuptools.setup(
|
|
packages=setuptools.find_packages(include=("deepinfer", "deepinfer.*", "deep_cluster", "deep_cluster.*")),
|
|
install_requires=get_requirements(),
|
|
ext_modules=[deepinfer_extension],
|
|
cmdclass={"build_ext": BuildExtension},
|
|
)
|