DeepInfer/docs/neuron/README.md

# 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`](../../deepinfer/neuron_ops/NEURON_PORT.md),
[`vllm_neuron_pr71_analysis.md`](vllm_neuron_pr71_analysis.md) (the vllm-neuron port this follows),
[`vllm_neuron_compile_notes.md`](vllm_neuron_compile_notes.md). Tiny-checkpoint correctness harness:
[`examples/neuron/README.md`](../../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

```bash
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).

<details><summary>中文</summary>

在一台 **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）。
</details>