DeepInfer/docs/neuron

DeepInfer on AWS Trainium (Neuron)

DeepSeek-V4.1-Flash-1017 on one trn2.48xlarge (16 Trainium2 devices, 64 NeuronCores at LNC=2) through native torch-neuronx (the PyTorch neuron device; no torch-xla). Basic mode: no DSpark/MTP, no zero-overhead scheduling, no prefix cache, text only.

Design notes: deepinfer/neuron_ops/NEURON_PORT.md, vllm_neuron_pr71_analysis.md (the vllm-neuron port this follows), vllm_neuron_compile_notes.md. Tiny-checkpoint correctness harness: examples/neuron/README.md.

How it runs

• Kernels. Pure-torch drop-ins with Ascend API semantics for deep_gemm / flash_mla / deep_ep / deep_select / tile_kernels (deepinfer/neuron_ops/vendor, tilelang_ops/neuron). The hot paths go through the NKI kernels of vllm-neuron#71 (deepinfer/neuron_ops/nki, Apache-2.0): block-scaled fp8/fp4 GEMMs on both physical cores, sparse MLA attention, indexer scores, compressor pooling, mHC, RMSNorm, SwiGLU. nkilib's rotational_topk and cascaded_max are used for sampling.
• Graphs. The step runs eagerly and compiles per-layer region graphs (neuron_ops/regions.py), shared by every layer of a kind. Prompt chunks run through the same graphs as decode-shaped steps padded to a token bucket. Context-length buckets size the dense compressed-KV scopes.
• fp8. trn2 has no OCP e4m3fn in graphs. The deepinfer_neuron torch.compile backend carries fp8 values as uint8 with arithmetic codecs, and checkpoint weights are rescaled into the e4m3 range (max 240) at load.
• Parallelism. EP=64 over all cores. --cp 64 makes the 64 ranks one TP group that runs one batch together: attention heads are sharded (one all-reduce a layer) and the MoE needs no all-gather. This is vllm-neuron#71's layout and the recommended setting for latency. --cp 1 is DP-attention.
• Engram tables sit in /dev/shm, shared by the ranks of the instance.

Run

git clone -b neuron https://github.com/whn09/DeepInfer.git && cd DeepInfer
bash tests/neuron/setup_trn2.sh          # venv on the DLAMI torch-neuronx venv, builds the C++ extension
export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0 DEEPINFER_ENFORCE_EAGER=0 \
       DEEPINFER_NEURON_PREFILL_BUCKET=128 DEEPINFER_MAX_BACKTRACKING_TOKENS=128 \
       DEEPINFER_NEURON_CTX_BUCKETS=2048,4096 DEEPINFER_USE_MULTIMODAL_POOL=0
# offline generation
python examples/neuron/offline_inference_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 \
    --ep 64 --cp 64 --max_model_len 8448 --num_kv_blocks 160 --max_num_batched_tokens 128 --max_tokens 64
# TTFT / TPOT matrix (vllm-neuron#71's method: random prompts, output 128, ignore_eos, greedy)
python examples/neuron/bench_neuron.py --model_name /path/to/DeepSeek-V4.1-Flash-1017 \
    --ep 64 --cp 64 --max_model_len 8448 --num_kv_blocks 160

The first run compiles the region graphs (~30-60 min, cached afterwards in the Neuron compile cache). An OpenAI-compatible server (deepseek-ai/DeepInfer#147, with a /v1/completions endpoint for vllm bench serve) is in serve/.

Performance (bs=1, greedy, output 128; whn09/DeepInfer#29)

input	TTFT	TPOT	vllm-neuron#71 TTFT / TPOT
512	5.0 s	447 ms	0.83 s / 117-127 ms
1024	9.0 s	447 ms	1.37 s
2048	17.1 s	456 ms	2.54 s
8192	66.5 s	471 ms	12.2 s

#71 decodes with DSpark (acceptance length ~2.75); its plain decode step is ~340 ms. Work in progress: prefill attention (#26), DSpark (#15), accuracy evaluation (#7).

中文

在一台 trn2.48xlarge（16 个 Trainium2 芯片，LNC=2 下 64 个 NeuronCore）上，用原生 torch-neuronx（PyTorch 的 neuron device，不依赖 torch-xla）运行 DeepSeek-V4.1-Flash-1017。目前是基础模式：没有 DSpark/MTP、没有零开销调度和 prefix cache，只支持纯文本。

• Kernel：deep_gemm / flash_mla / deep_ep / deep_select / tile_kernels 都有遵循 Ascend API 语义的纯 torch 替代实现；热路径走 vllm-neuron#71 的 NKI kernel（Apache-2.0）：块标度 fp8/fp4 GEMM（用满两个物理核）、稀疏 MLA attention、indexer、压缩池化、mHC、RMSNorm、SwiGLU；采样用 nkilib 的 rotational_topk 和 cascaded_max。
• 图：每个 step 以 eager 方式运行，每层编译按层的区域图，同类层共享同一张图；prompt chunk 按 token bucket pad 成 decode 形状后，复用同一套图；dense 压缩 KV 的宽度按上下文长度分 bucket。
• fp8：trn2 的图里不支持 OCP e4m3fn。deepinfer_neuron 这个 torch.compile backend 用 uint8 承载 fp8 值，并用算术方式编解码；加载时把权重缩放到 e4m3 的范围（最大 240）内。
• 并行：EP=64 覆盖全部核。--cp 64 让 64 个 rank 组成一个 TP 组，一起跑同一个 batch：attention 按头切分（每层一次 all-reduce），MoE 不再需要 all-gather。这与 vllm-neuron#71 的布局相同，是低延迟场景的推荐配置。--cp 1 则是 DP-attention。
• Engram 表放在 /dev/shm，同一台实例上的所有 rank 共享。

运行方式和性能数据见上文英文部分。首次运行要编译区域图（约 30–60 分钟），之后会命中 Neuron 编译缓存。serve/ 下有 OpenAI 兼容的 server（来自 deepseek-ai/DeepInfer#147，补了供 vllm bench serve 使用的 /v1/completions）。#71 开了 DSpark（接受长度约 2.75），不开时单个 decode step 约 340 ms。仍在进行中的工作：prefill attention（#26）、DSpark（#15）、精度评测（#7）。