- 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>
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>