coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-11-09 14:49:45 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/neuron_ops/NEURON_PORT.md
whn09 ed49792a0c Add AWS Trainium (Neuron) platform, basic mode — part 1 
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>
2026-10-20 18:34:51 +08:00

3.8 KiB
Raw Blame History

DeepInfer on AWS Trainium (native torch-neuronx, no XLA)

Stack

  • torch 2.14 + torch-neuronx 2.14.3 (PrivateUse1 device neuron, eager + torch.compile(backend="neuron")), Neuron SDK 2.33 (runtime 2.35). No torch-xla, no libtorch-neuronx-lite.
  • Target: trn2.48xlarge (64 logical NeuronCores, LNC=2, 24 GB HBM per core), Neuron DLAMI.
  • Python env: a venv layered on the DLAMI's /opt/aws_neuronx_venv_torch_neuronx_2_14_0 (add its bin/ to PATH so neuronx-cc is found), plus requirements.txt minus tilelang / tile-kernels / torch.
  • Build: DEEPINFER_ACCELERATOR=neuron ./scripts/build.sh build (host-only C++ extension, -DWITH_NEURON).

Platform design

  • PlatformNeuron (deepinfer/platform.py): is_npu is true (Neuron follows the Ascend kernel layouts / code paths), is_neuron is true, tilelang_ops_backend = "neuron".
  • deepinfer/neuron_ops/transfer_to_neuron.py maps "cuda" devices and torch.cuda.* onto the target device. DEEPINFER_NEURON_DEVICE=neuron|cpu selects the target; cpu runs everything on the host.
  • deepinfer/neuron_ops/vendor/ is put at the front of sys.path (deepinfer/init.py) so import deep_gemm, flash_mla, deep_ep, deep_select, tile_kernels, tilelang resolve to the pure-torch packages there.
  • deepinfer/model_executor/layers/tilelang_ops/neuron/ is the kernel backend: every xxx_tl factory that asc/ provides, same name and signature, returning a callable with the same call signature and the same in-place output semantics.

Scope: basic adaptation only

Optional optimizations stay OFF on Neuron (defaults set in deepinfer/init.py): CUDA graphs (eager), DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0), zero-overhead scheduling, MTP varlen, EPLB, CP>1, YOCO bounded replay, batch-invariant mode, GPU-only mega kernels, VL/vision. Kernels reachable only from those features: export the name with the right signature but raise NotImplementedError("<name>: not supported on Neuron (basic mode)"). Fully implement what the model architecture needs (MLA attention, DSA indexer/top-k, compressor, mHC, Engram, MoE with FP8/FP4 dequant, norms/rope, KV cache store/load, basic sampling).

Rules for kernel code (vendor packages and tilelang_ops/neuron)

  1. Pure torch, device-agnostic: allocate on the input tensors' device; never hardcode cuda/npu/neuron. Must run with tensors on CPU (correctness now) and on neuron (later).
  2. Semantics follow the Ascend variant (DeepGEMM-Ascend, DeepEP-Ascend, tilelang_ops/asc/, the PLATFORM.is_npu branches); where Ascend has no variant, follow the CUDA one.
  3. Output semantics must match exactly: if a kernel writes into caller-provided tensors, write in place (out.copy_(...), indexed assignment); keep the caller-visible shape/dtype/stride contract.
  4. Prefer static shapes and no host syncs (.item(), .tolist(), nonzero, boolean-mask indexing) in per-step model code — the decode step will later be compiled to one static graph per batch bucket. Where a data-dependent shape is unavoidable, isolate it and leave a # NEURON: host sync comment.
  5. Neuron dtype notes: prefer int32 over int64 for indices on device; fp8 is torch.float8_e4m3fn; FP4 (e2m1) is packed two per byte — unpack with bit ops; ue8m0 scales are exponent bytes (value = 2**(b-127)).
  6. Correctness over speed. Match the reference numerics (fp32 accumulation where the kernel accumulates in fp32).
  7. Each file has a short module docstring saying which upstream API it mirrors.

Testing

CPU mode (same code, tensors on the host):

export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu

Compare against the upstream references (*_ref functions next to each tilelang_ops interface, tile_kernels.torch, and the reference implementations inside vendor/<Repo>/tests).