Neuron packaging: requirements-neuron.txt, build.sh support, Dockerfile, README

- 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>
2026-10-28 14:57:09 +08:00 · 2026-10-21 14:57:51 +08:00 · 2026-10-21 14:57:51 +08:00 · d0f7e4f69c
commit d0f7e4f69c
parent cc90926a1c
7 changed files with 174 additions and 6 deletions
--- a/README.md
+++ b/README.md
@ -34,6 +34,7 @@ Supported hardware:

 - NVIDIA SM100 GPUs
 - Huawei Ascend 950 NPUs
+- AWS Trainium2 (native torch-neuronx, basic mode — see [DeepInfer on AWS Trainium](#deepinfer-on-aws-trainium))

 Required for a complete inference service (not included in this release):

@ -50,13 +51,14 @@ Required for a complete inference service (not included in this release):
 - Python 3.12
 - CUDA 13 (NVIDIA only)
 - CANN 9.2 (Ascend only)
+- Neuron SDK 2.33 with torch-neuronx 2.14 (Trainium only)

 ### Build from source

 Before building DeepInfer, please create and activate an empty Python 3.12 environment.

 ``` bash
-# Use `ascend` instead for the Ascend platform
+# Use `ascend` instead for the Ascend platform, `neuron` for AWS Trainium
 # Keep it exported in the environment where you run DeepInfer
 export DEEPINFER_ACCELERATOR=nvidia
 ./build.sh
@ -74,6 +76,74 @@ python examples/offline_inference.py

 See [examples](examples/README.md) for details.

+## DeepInfer on AWS Trainium
+
+DeepInfer runs on Trainium2 through the native PyTorch `neuron` device (torch-neuronx 2.14, no torch-xla),
+structured like the Ascend port: `PlatformNeuron`, the `tilelang_ops/neuron` kernel backend and pure-torch
+drop-ins for DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels in `deepinfer/neuron_ops/vendor`.
+Design notes and conventions: [deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md).
+
+**Status — basic mode.** Eager execution; DSpark/MTP, zero-overhead scheduling, CP, EPLB, YOCO bounded replay and
+VL are off. EP > 1 runs on real neuron collectives. A tiny checkpoint in the real FP8/FP4 layout matches the
+official reference implementation within its own noise floor, on CPU and on a NeuronCore
+([examples/neuron](examples/neuron/README.md)). Performance work (static graphs, NKI kernels) and the remaining
+features are tracked in the [roadmap](https://github.com/whn09/DeepInfer/issues/19).
+
+On a Neuron DLAMI (trn2.48xlarge, SDK 2.33):
+
+``` bash
+bash tests/neuron/setup_trn2.sh          # venv on the DLAMI torch-neuronx env, build, tiny checkpoint + golden
+bash tests/neuron/run_all.sh cpu         # host-only tests
+CORES=0-7 bash tests/neuron/run_all.sh device
+```
+
+or build by hand in a venv layered on the DLAMI's torch-neuronx environment (the `torch-neuronx` wheel on the
+Neuron pip repository needs glibc >= 2.43, newer than Ubuntu 24.04's; the DLAMI ships a compatible build):
+
+``` bash
+python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \
+  > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate
+export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin   # neuronx-cc
+export DEEPINFER_ACCELERATOR=neuron
+./build.sh                               # requirements-neuron.txt (torch / torch-neuronx already satisfied) + C++ build
+python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64
+```
+
+<details>
+<summary>中文</summary>
+
+DeepInfer 通过原生 PyTorch 的 `neuron` 设备（torch-neuronx 2.14，不依赖 torch-xla）在 Trainium2 上运行，结构与
+Ascend 移植一致：`PlatformNeuron`、`tilelang_ops/neuron` kernel 后端，以及 `deepinfer/neuron_ops/vendor` 中
+DeepGEMM / FlashMLA / DeepEP / DeepSelect / TileKernels 的纯 torch 实现。设计说明与约定见
+[deepinfer/neuron_ops/NEURON_PORT.md](deepinfer/neuron_ops/NEURON_PORT.md)。
+
+**现状：basic mode。** 逐算子即时执行；DSpark/MTP、零开销调度、CP、EPLB、YOCO 有界重放和 VL 均关闭。EP > 1 使用
+真实的 neuron 集合通信。按真实 FP8/FP4 布局生成的小 checkpoint 在 CPU 和 NeuronCore 上都与官方参考实现一致（误差在
+参考实现自身的噪声范围内，见 [examples/neuron](examples/neuron/README.md)）。性能优化（静态图、NKI kernel）和其余特性
+在 [roadmap](https://github.com/whn09/DeepInfer/issues/19) 中跟踪。
+
+在 Neuron DLAMI（trn2.48xlarge，SDK 2.33）上：
+
+``` bash
+bash tests/neuron/setup_trn2.sh          # 基于 DLAMI 的 torch-neuronx 环境建 venv、编译、生成小 checkpoint 和真值
+bash tests/neuron/run_all.sh cpu         # 只用 host 的测试
+CORES=0-7 bash tests/neuron/run_all.sh device
+```
+
+或者在叠加于 DLAMI torch-neuronx 环境之上的 venv 里手动构建（Neuron pip 源上的 `torch-neuronx` wheel 需要
+glibc >= 2.43，比 Ubuntu 24.04 的版本新；DLAMI 自带的是兼容的构建）：
+
+``` bash
+python3.12 -m venv venv && echo "import site; site.addsitedir('/opt/aws_neuronx_venv_torch_neuronx_2_14_0/lib/python3.12/site-packages')" \
+  > venv/lib/python3.12/site-packages/zz_neuron.pth && source venv/bin/activate
+export PATH=$PATH:/opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin   # neuronx-cc
+export DEEPINFER_ACCELERATOR=neuron
+./build.sh                               # requirements-neuron.txt（torch / torch-neuronx 已满足）+ C++ 编译
+python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 --ep 64
+```
+
+</details>
+
 ## Performance benchmarks

 | Prefill | Decode |
--- a/deepinfer/platform.py
+++ b/deepinfer/platform.py
@ -285,6 +285,7 @@ class PlatformNeuron(Platform):
 def get_platform() -> Platform:
    from deep_cluster.hardware import accelerator_kind

-    kind = accelerator_kind()
+    # An explicitly exported DEEPINFER_ACCELERATOR wins (e.g. Neuron CPU mode on a host without /dev/neuron*).
+    kind = os.environ.get("DEEPINFER_ACCELERATOR") or accelerator_kind()
    platform_class = {"ascend": PlatformNPU, "neuron": PlatformNeuron}.get(kind, PlatformGPU)
    return platform_class()
--- a/docker/neuron/Dockerfile
+++ b/docker/neuron/Dockerfile
@ -0,0 +1,31 @@
+# DeepInfer on AWS Trainium (Neuron basic mode), on top of the native PyTorch Neuron DLC.
+#
+# The base image ships torch 2.14 + torch-neuronx 2.14 (no torch-xla), neuronx-cc, nki and the Neuron tools on
+# Ubuntu 26.04 / Python 3.14 (Neuron DLCs are Ubuntu 26.04-only from SDK 2.33; the torch-neuronx wheel needs
+# glibc >= 2.43). The Neuron driver (aws-neuronx-dkms) must be installed on the host.
+#
+#   docker build -f docker/neuron/Dockerfile -t deepinfer-neuron .
+#   docker run --rm -it --privileged --ipc=host --network=host \
+#       -v /mnt/nvme/models:/models deepinfer-neuron \
+#       python examples/neuron/offline_inference_neuron.py --model_name /models/DeepSeek-V4.1-Flash-1017 --ep 64
+ARG BASE_IMAGE=public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-sdk2.33.0-ubuntu26.04
+FROM ${BASE_IMAGE}
+
+ENV DEEPINFER_ACCELERATOR=neuron \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_ROOT_USER_ACTION=ignore
+
+WORKDIR /opt/deepinfer
+COPY . .
+
+# requirements-neuron.txt (torch / torch-neuronx are already satisfied by the base image), no vendor build
+# (Neuron uses the pure-torch drop-ins in deepinfer/neuron_ops/vendor), then the host-only C++ extension.
+# No NeuronCore is visible during `docker build`: keep torch from auto-loading torch_neuronx (it queries the device
+# when it initializes); compiling the extension only needs the torch headers. The self-check runs in Neuron CPU mode.
+RUN TORCH_DEVICE_BACKEND_AUTOLOAD=0 MAX_JOBS=16 ./build.sh \
+ && TORCH_DEVICE_BACKEND_AUTOLOAD=0 DEEPINFER_NEURON_DEVICE=cpu python -c "import deepinfer, torch; deepinfer.PLATFORM; \
+from deepinfer._C import engine; from deepinfer.model_executor.layers.tilelang_ops import ops; \
+print('deepinfer-neuron OK:', type(deepinfer.PLATFORM).__name__, 'torch', torch.__version__)"
+
+ENV PYTHONPATH=/opt/deepinfer
+CMD ["/bin/bash"]
--- a/docker/neuron/Dockerfile.dockerignore
+++ b/docker/neuron/Dockerfile.dockerignore
@ -0,0 +1,8 @@
+# Applies to docker/neuron/Dockerfile only (BuildKit per-Dockerfile ignore file).
+.git
+build
+deepinfer/_so
+**/__pycache__
+*.pyc
+# Kernel sources for CUDA / Ascend; Neuron images use the vendored pure-torch packages instead.
+vendor
--- a/requirements-neuron.txt
+++ b/requirements-neuron.txt
@ -0,0 +1,46 @@
+# AWS Trainium (Neuron SDK 2.33+, native torch-neuronx -- no torch-xla).
+# Same dependencies as requirements.txt + requirements-dev.txt, except:
+#   * tilelang / tile-kernels: not needed -- Neuron uses the pure-torch drop-ins vendored in
+#     deepinfer/neuron_ops/vendor (deep_gemm, flash_mla, deep_ep, deep_select, tile_kernels, tilelang);
+#   * torch comes with torch-neuronx (torch 2.14), not the CUDA/NPU builds.
+# On a Neuron DLAMI the torch-neuronx stack is already in /opt/aws_neuronx_venv_torch_neuronx_2_14_0;
+# tests/neuron/setup_trn2.sh layers a venv on it instead of reinstalling.
+--extra-index-url https://pip.repos.neuron.amazonaws.com
+--find-links https://download.pytorch.org/whl/cpu/torch/
+torch==2.14.0
+torch-neuronx==2.14.*
+neuronx-cc==2.*
+nki>=0.6
+
+einops==0.8.2
+huggingface_hub==0.36.2
+ninja==1.13.0
+numpy==2.3.5
+orjson==3.11.7
+pillow==11.3.0
+posix_ipc==1.3.2
+psutil==7.2.2
+pybind11==3.0.1
+pyzmq==27.1.0
+requests==2.32.5
+safetensors==0.8.0
+scipy==1.17.0
+sympy==1.14.0
+tabulate==0.9.0
+tokenizers==0.21.4
+tqdm==4.67.3
+transformers==4.50.0
+xgrammar==0.2.2
+
+# development / tests
+clang-format==22.1.3
+execnet==2.1.2
+pybind11-stubgen==2.5.5
+pytest==8.4.2
+pytest-xdist==3.8.0
+ruff==0.15.0
+tree-sitter==0.26.0
+tree-sitter-cpp==0.23.4
+matplotlib==3.11.2
+fastapi==0.128.0
+uvicorn==0.40.0
--- a/scripts/build_vendor.sh
+++ b/scripts/build_vendor.sh
@ -15,8 +15,13 @@ case "${DEEPINFER_ACCELERATOR:-nvidia}" in
    PACKAGES=(DeepSelect FlashMLA DeepEP-Ascend DeepGEMM-Ascend)
    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
    ;;
+  neuron)
+    # Neuron uses the pure-torch drop-ins vendored in deepinfer/neuron_ops/vendor; nothing to build.
+    echo "✅ Neuron: vendored kernel packages are pure torch (deepinfer/neuron_ops/vendor), nothing to build"
+    exit 0
+    ;;
  *)
-    echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia or ascend)" >&2
+    echo "Unsupported DEEPINFER_ACCELERATOR: '${DEEPINFER_ACCELERATOR}' (expect nvidia, ascend or neuron)" >&2
    exit 1
    ;;
 esac
--- a/setup.py
+++ b/setup.py
@ -9,8 +9,11 @@ from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CppExtension, C

 from deep_cluster.hardware import accelerator_kind

-IS_NPU = accelerator_kind() == "ascend"
-IS_NEURON = accelerator_kind() == "neuron"
+# 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__))

@ -96,7 +99,11 @@ 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]
+        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(