From 2310c2db8ac3dd028f1f9309fa6c2dc7fe11310b Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 08:31:57 +0800 Subject: [PATCH] docs/neuron/README.md: running DeepSeek-V4.1-Flash on trn2, performance Co-Authored-By: Claude Opus 5.5 --- docs/neuron/README.md | 72 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 docs/neuron/README.md diff --git a/docs/neuron/README.md b/docs/neuron/README.md new file mode 100644 index 0000000..5f608af --- /dev/null +++ b/docs/neuron/README.md @@ -0,0 +1,72 @@ +# 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). + +
中文 + +在一台 **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）。 +