mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
d0f7e4f69c
- requirements-neuron.txt: requirements.txt + dev deps without tilelang / tile-kernels (pure-torch drop-ins are vendored), torch 2.14 + torch-neuronx 2.14 from the Neuron repository. - build_vendor.sh: nothing to build for neuron; setup.py drops tilelang / tile-kernels from install_requires on neuron. - setup.py and get_platform(): an exported DEEPINFER_ACCELERATOR wins over device detection, so images build where no accelerator is visible. - docker/neuron/Dockerfile (+ per-Dockerfile ignore): on the native PyTorch Neuron DLC public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314- sdk2.33.0-ubuntu26.04; build.sh with TORCH_DEVICE_BACKEND_AUTOLOAD=0 and an import self-check in Neuron CPU mode. - README: Trainium in supported hardware and a "DeepInfer on AWS Trainium" section (English + folded Chinese); the pip torch-neuronx wheel needs glibc >= 2.43, so on Ubuntu 24.04 layer the venv on the DLAMI env. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
115 lines
4.0 KiB
Python
115 lines
4.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, CppExtension, CUDAExtension
|
|
|
|
from deep_cluster.hardware import accelerator_kind
|
|
|
|
# An explicitly exported DEEPINFER_ACCELERATOR (build.sh requires it) wins over device detection, so that images
|
|
# can be built where the accelerator is not visible (e.g. `docker build`).
|
|
_ACCELERATOR = os.environ.get("DEEPINFER_ACCELERATOR") or accelerator_kind()
|
|
IS_NPU = _ACCELERATOR == "ascend"
|
|
IS_NEURON = _ACCELERATOR == "neuron"
|
|
|
|
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,
|
|
)
|
|
elif IS_NEURON:
|
|
# Native torch-neuronx device: host-only C++, no device toolkit headers needed.
|
|
deepinfer_extension = CppExtension(
|
|
name="deepinfer._C",
|
|
sources=[
|
|
*COMMON_SOURCES,
|
|
"csrc/neuron_utils.cpp",
|
|
"csrc/graph_timer/graph_timer_neuron.cpp",
|
|
],
|
|
extra_compile_args={
|
|
"cxx": [*CXX_FLAGS, "-fvisibility=hidden", "-DWITH_NEURON"],
|
|
},
|
|
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)
|
|
reqs = [line for line in lines if line]
|
|
if IS_NEURON:
|
|
# Neuron vendors pure-torch tilelang / tile_kernels (deepinfer/neuron_ops/vendor); see requirements-neuron.txt.
|
|
reqs = [r for r in reqs if not r.startswith(("tilelang", "tile-kernels"))]
|
|
return reqs
|
|
|
|
|
|
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},
|
|
)
|