mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
ed49792a0c
Run DeepInfer on the native torch-neuronx device (PrivateUse1 `neuron`, torch-neuronx 2.14, no torch-xla), following the Ascend port's structure. - deep_cluster: Neuron discovery (one DeepInfer device = one logical NeuronCore, NEURON_RT_VISIBLE_CORES, NUMA from neuron-ls). - PlatformNeuron: takes the Ascend (is_npu) kernel layouts, adds is_neuron; DEEPINFER_NEURON_DEVICE=cpu runs the same code on the host for correctness work. - neuron_ops/transfer_to_neuron: map "cuda" devices, torch.device and torch.cuda.* onto neuron:0 (counterpart of torch_npu transfer_to_npu). - C++ extension: -DWITH_NEURON host-only build; torch-allocated buffers and tensor-view slot copies instead of raw device-pointer copies. - Pure-torch drop-in vendor packages (Ascend API semantics): deep_gemm, flash_mla, deep_select, tile_kernels, tilelang (constants only). - tilelang_ops/neuron backend: KV cache / metadata / norm / compress kernels; CP, ZOS and VL entry points raise NotImplementedError. - Basic mode defaults on Neuron: eager, no DSpark/MTP drafting, no zero-overhead scheduling; MoE avoids the GPU-only mega kernel; mHC prenorm GEMM uses the Ascend single-split reduction. Not yet included: deep_ep, the sampling kernels, end-to-end harness. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
108 lines
3.6 KiB
Python
108 lines
3.6 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
|
|
|
|
IS_NPU = accelerator_kind() == "ascend"
|
|
IS_NEURON = accelerator_kind() == "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)
|
|
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},
|
|
)
|