mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
02e099fff8
- NeuronGraphRunner: the CudaGraphRunner contract over torch.compile(dynamic=False), one graph per batch size (DEEPINFER_ENFORCE_EAGER=0, backend DEEPINFER_NEURON_COMPILE_BACKEND). Layer cache views are re-sliced from their root allocations inside the graph, so AOT never regenerates aliased mutated inputs with as_strided. - Decode MoE with fixed shapes: packed all-gather, dense local experts, fp32 reduce-scatter (neuron_ops/static_moe.py, deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a). - Graph-break and compile fixes: hide torch.cuda during compile, torch.device equality for FX codegen, hash-based force_random routing, Engram hash/fetch on the host (exact int64), env vars resolved before tracing, Movable annotation, no .item()/nonzero/as_strided writes on the decode path, KV stores route padding slots to a scratch block kept past the pool, native-dtype gather under compile. - fp8: trn2 reads float8_e4m3fn as IEEE e4m3 (max 240). Activations/KV quantize to 224 (DEEPINFER_NEURON_FP8_MAX), checkpoint fp8 weights are rescaled at load (bytes / 2, ue8m0 + 1), and neuronx-cc gets --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 via --internal-hlo2tensorizer-options. The uint8-carrier backend (neuron_ops/neuron_compile.py, "deepinfer_neuron") stays as a fallback. - Tests: tests/neuron/compile; run_all.sh runs kernel unit tests at the e4m3fn range (448) and the end-to-end stages at 224. Docs: vllm-neuron PR #71 analysis and compile notes. Host: tiny model 9/9 vs golden (all prompts) eager, compiled path (aot_eager) matches eager; CPU suite green. Device: the whole decode graph still crashes neuronx-cc (NCC_IINL902, being bisected). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
282 lines
33 KiB
Markdown
282 lines
33 KiB
Markdown
# How vllm-neuron compiles a decode step (notes for the DeepInfer Neuron port)
|
||
|
||
Source: local checkout `Trn2/vllm-neuron` (HEAD `f32079d`, "Merge pull request #63 add-contrib-folder"). Paths are relative to the repo root. Read-only survey. Nothing in vllm-neuron was changed.
|
||
|
||
**Caveat.** vllm-neuron does **not** use `torch-neuronx` or its `"neuron"` backend. It compiles through `libtorch-neuronx-lite` (`requirements/core.txt:7`), which carries its own copy of the native lowering: Dynamo, then AOTAutograd with Neuron decomps, then torch-mlir, StableHLO and neuronx-cc (`docs/design/compilation/how_vllm_is_integrated_with_torchneuron.md`). The lite source is not in this checkout. Statements about what lite does internally come from comments and docs, not from code.
|
||
|
||
## 1. Compile backend and compiler flags
|
||
|
||
- `envs.get_compile_backend_name()` (`vllm_neuron/envs.py:789-803`) has two routes:
|
||
- **Native route** (the default, `VLLM_NEURON_BACKEND` unset, `neuron_native` or `neuron`): it calls `libtorch_neuronx_lite.compile.native_backend.register()` and returns the name that call registers. That name is deliberately **not** `"neuron"`, which "stays reserved for the torch_neuronx package".
|
||
- **XLA route** (deprecated) and CPU mode: it returns `"neuron_libtorch"`, registered at import (`vllm_neuron/__init__.py:~128`).
|
||
- `torch-neuronx` and lite cannot coexist unless `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. Otherwise you get "multiple backend fallbacks for PrivateUse1".
|
||
- The torch.compile call is in `vllm_neuron/vllm/worker/neuron_model_runner_v2.py:1903-1908`:
|
||
`torch.compile(NeuronForwardAdapter(model), backend=envs.get_compile_backend_name(), fullgraph=not DEBUG, options=self.compile_options)`.
|
||
- `torch._dynamo.config.cache_size_limit = 2**62` (`:1860`). Buckets are meant to specialize; recompiles are caught by the stance instead.
|
||
- The serve-time forward runs under `torch.compiler.set_stance("fail_on_recompile")` (`utils/neuron_utils.py:26-32`, used at `_serve_forward_guard`, `neuron_model_runner_v2.py:1223`).
|
||
- Optional `guard_filter_fn` (`envs.py:814`) for warm start: `keep_portable_guards_unsafe` or `skip_all_guards_unsafe`.
|
||
- `compile_options` (`neuron_model_runner_v2.py:387`, `:672-701`):
|
||
```
|
||
{"alias_meta_to_neuron": True, # lite: thread aliased (mutated) outputs back onto inputs
|
||
"compiler_args": ["--auto-cast=none", "--verbose=35", f"-O{vllm optimization_level}",
|
||
"--internal-hlo2tensorizer-options=--modular-flow-mac-threshold=10"
|
||
+ (" --experimental-unsafe-fp8e4m3fn-as-fp8e4m3" if fp8 and target not trn3),
|
||
"--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop",
|
||
(+ explicit --lnc=N only when it differs from the compiler default)]}
|
||
```
|
||
- `--enable-nested-dynamic-loop` is needed because NKI kernels inlined into the parent BIR have nested `scf.for`, and birverifier rejects the extra back-edges (`neuron_model_runner.py:2071-2076`).
|
||
- All hlo2tensorizer options must go into **one** string. A second `--internal-hlo2tensorizer-options` replaces the first (`neuron_model_runner.py:2037-2047`).
|
||
- The V1 runner adds `--enable-annotate-fsdp-cc=false` and `--model-type=generic` for the GPT-OSS giga-kernel graph (`:2047`, `:2079`). torch-neuronx injects `--model-type transformer`, and neuronx-cc keeps the last flag it sees.
|
||
- **Per-graph args.** A model may define `compiler_args_for(graph_type, example_inputs, base)` (`neuron_model_runner_v2.py:703-723`; `docs/design/framework/per_bucket_compiler_args.md`). Qwen3-VL and ASR compile the decode and encoder graphs at **-O3**, because "-O3 disables neuronx-cc modular flow; the -O1 default adds a PSEUDO_FUNCTION_RETURN/EVENT_SEMAPHORE wrapper on each collective" (`model/qwen3_vl/model_common.py:290-301`). Prefill keeps -O2.
|
||
- The native route defaults to O2 and the XLA route to O1 (`docs/guides/how-to-choose-compilation-route.md`).
|
||
|
||
## 2. KV cache as graph input (the synthetic-base problem)
|
||
|
||
- **The KV cache is not a forward argument.** `bind_kv_cache` stores plain tensor attributes on each attention module: `layer.self_attn.k_cache = kv[name][0]`, `.v_cache = kv[name][1]` (`model/qwen3/model.py:1160-1166`). Dynamo lifts these attributes as graph inputs. The adapter signature (`worker/neuron_forward_adapter.py:209-227`) carries only metadata and state tensors.
|
||
- **Allocation** (V2, `neuron_model_runner_v2.py:917-1215`):
|
||
- One `int8` backing tensor is created (`:1140`).
|
||
- For each layer it takes `raw = backing[start:start+span]`, then `raw.view(dtype).view(2, num_blocks, kv_heads, block_size, head_size)`. **K = typed[0] and V = typed[1]**: two non-overlapping views of one per-layer region.
|
||
- **The trick that avoids the synthetic base.** View objects are memoized: `shared_dtype_view` and `shared_kv_slices` are keyed by `(id(raw), dtype, kv_shape)` (`:1148-1186`). Layers that share bytes then receive the **exact same tensor objects**, and Dynamo dedups them into one placeholder. The comment at `:1160-1166` says: "Re-slicing shared storage per layer gives Dynamo overlapping inputs that AOTAutograd merges into a whole-buffer synthetic base, widening a slot-sized KV update into a full-buffer write."
|
||
- In short, the views AOTAutograd sees never partially overlap: K and V are disjoint, and layers are disjoint. AOTAutograd's overlap check then keeps them as independent mutated inputs and creates no synthetic base and no `as_strided` replay.
|
||
- A host preflight rejects pooled layers whose `(dtype, heads, block, head_size)` differ, because they would produce distinct, overlapping view objects (`:1086-1110`).
|
||
- **2**32 offset hazard.** The V1 runner deliberately allocates **one `torch.zeros` per layer** rather than slicing one giant buffer. Its reason: "views into a multi-GiB shared storage carry element offsets past 2**32, which the lowering path can truncate (silent cross-layer KV corruption)" (`neuron_model_runner.py:10988-10995`). V2 slices a single backing. **For DeepInfer, per-layer allocations are the safe choice.**
|
||
- **KV writes are plain `index_put_` inside the graph**, not a custom op and not a kernel-side cache update:
|
||
- `paged_slot_indices(slot_mapping, block_size, num_blocks)` (`model/kv_cache.py:80`) clamps invalid or `-1` slots to null block 0 **before** `//` and `%`. A `-1` slot becomes an out-of-bounds DGE scatter, which wedges the NEFF execution queue.
|
||
- Then `k_cache.index_put_((block_idx, head_idx, pos_idx), k_flat)` (`model/qwen3/model.py:336-343` for prefill, `:582-589` for decode).
|
||
- The decode NKI kernel is called with `update_cache=False` and returns `K_new` and `V_new`. The model then scatters them itself (`:525-589`, `qwen3_vl/utils/decode_kv.py:60` `scatter_new_kv`).
|
||
- **Mutation becomes aliasing.**
|
||
- Lite's `alias_meta_to_neuron=True` turns every mutated input into an extra output plus an `input_output_alias` (donation). The XLA-route equivalents are documented in `docs/design/compilation/aliasing_output_rewrite_pass.md` (mutated inputs are appended to the outputs, giving `io_map{out_idx: in_idx}`) and `inplace_to_outofplace_pass.md`. NKI kernels declare `operand_output_aliases` for the same purpose.
|
||
- vllm-neuron uses **no** `mark_static_address`, **no** `keep_inference_input_mutations` and no custom `torch._functorch` config.
|
||
- **Runner state follows the same pattern.** `token_ids.index_put_(...)` runs in-graph (`neuron_forward_adapter.py:190`), and the aliased bundle is returned.
|
||
- The next decode reads a **clone** of the sampled store, a non-aliased output (`:345-348`). That keeps TP sampling collectives ordered across ranks.
|
||
- The output future is kept alive until the next step (`neuron_model_runner_v2.py:1453-1460`). Otherwise in-flight NRT input slices get freed and TP hangs.
|
||
|
||
## 3. Graph granularity
|
||
|
||
- **One graph = the whole forward**, for each (phase, bucket). It covers the input-id build, all decoder layers (a Python loop unrolled by Dynamo, `model/qwen3/model.py:896`), the LM head, on-device sampling and the state commit (`neuron_forward_adapter.py:272-356`).
|
||
- Prefill and decode are separate graphs, selected by whether `linear_idx` is None.
|
||
- NEFF count is `|num_seqs_buckets| × (|decode_ctx_buckets|+1)` (`docs/design/vllm/decode-context-length-bucketing.md:111-116`).
|
||
- **There is no per-layer compiled region**, and no mechanism to "compile one layer once and reuse it". Reuse of repeated layers is left to **neuronx-cc modular flow** (`--modular-flow-mac-threshold=10` at O1/O2). It is turned **off** with `-O3` for the latency-critical decode graph.
|
||
- The giga-kernel path goes further: one NKI kernel for all 36 GPT-OSS layers in a single decode NEFF, with weights stacked on a `[num_layers, ...]` axis (`docs/design/vllm/gpt-oss-gk-decode.md`).
|
||
- The ~120-NEFF design in the DeepSeek-V4.1 PR #71 (per-expert graphs, `@torch._dynamo.disable` MoE, host sync per layer) is the opposite of what the rest of the repo does. See `vllm_neuron_pr71_analysis.md`.
|
||
- Compile time is handled by parallel trace and compile across the TP group (`parallel_compile`, `neuron_model_runner_v2.py:2035-2056`), the compile cache, and the SPMD single-NEFF guard. The graph is not split.
|
||
|
||
## 4. Buckets and static shapes
|
||
|
||
- Defaults (`utils/bucket_utils.py:36-126`):
|
||
- Prefill token buckets are powers of 2 from 128 up to `max_num_batched_tokens`.
|
||
- Decode batch buckets are powers of 2 from 1 up to `max_num_seqs`.
|
||
- Optional KV-segment buckets (prefix caching) and decode context-length buckets.
|
||
- `_select_bucket` picks the smallest bucket that fits (`neuron_model_runner_v2.py:903`).
|
||
- Padding rules:
|
||
- Padded tokens use `PAD_TOKEN=0`. Padded slots use `slot=-1`, which `paged_slot_indices` redirects to null block 0.
|
||
- The block table uses the `-1` sentinel past each row's valid blocks, and the attention kernel skips OOB.
|
||
- Padded decode rows have `row_to_slot=-1` (`worker/neuron_input_build.py:1-13`; `docs/design/vllm/padding-batching-block-kv.md`).
|
||
- Metadata is built on the host with NumPy only and never reads device state. It is copied with `torch.from_numpy(...).to(device)` **outside** the graph each step (`neuron_model_runner_v2.py:2687-2712`).
|
||
- Python ints (`max_query_len`, `block_size`, `kv_segment_size`) are part of the metadata dict on purpose, so they specialize the graph.
|
||
- Decode `input_ids` are built **in-graph** from the previous step's on-device sampled tokens, so steady-state decode does no D2H.
|
||
- **No `.item()` in the graph.** Host reads happen only at output time (`.cpu().tolist()` of the sampled future).
|
||
- Scale floats are kept as Python floats next to the tensors, because `tensor.item()` causes a graph break (`model/llama3/model.py:443`).
|
||
|
||
## 5. Collectives inside graphs
|
||
|
||
- The world and model-parallel process groups are **Gloo**. Collectives are traced into the NEFF and never dispatched at runtime (`envs.py:858-866`).
|
||
- APIs used:
|
||
- `torch.distributed._functional_collectives.all_reduce(x, reduceOp="sum", group=...)`. The reduceOp must be a **string**, because a `ReduceOp` object breaks precompile pickling (`nn/embedding.py:256-260`, `nn/rpl.py:147-150`).
|
||
- `all_gather_tensor`.
|
||
- `torch.ops._c10d_functional.all_gather_into_tensor`, `all_to_all_single` and `wait_tensor` (`functional/collectives/all_to_all.py:127-193`).
|
||
- Plain in-place `dist.all_reduce` (through vLLM's `GroupCoordinator`) is also used. Dynamo remaps it to the functional form.
|
||
- Restrictions:
|
||
- **trn2: all_to_all is emulated** as all_gather plus a local slice, because "ENC lacks Mesh" (`all_to_all.py:10-33`).
|
||
- Every rank must trace byte-identical HLO (SPMD). Use a runtime `rank` tensor, not Python ints, and replace rank-dependent slices with `index_select(arange+offset)`. Each group registers its full replica-group mesh (`docs/design/compilation/spmd_compilation.md`, `patches/replica_mesh_patch.py`, `parallel/neuron_group_coordinator.py:122-130`).
|
||
- **A collective cannot read a NEFF I/O tensor** (`NCC_INLA001`). Copy the input to an intermediate first (`functional/collectives/reduce_scatter_v.py:80-95, 318-330`).
|
||
|
||
## 6. Known neuronx-cc workarounds found
|
||
|
||
| Symptom | Workaround | Where |
|
||
|---|---|---|
|
||
| birverifier rejects nested loops from inlined NKI | `--enable-verifier=false --enable-nested-dynamic-loop` | `neuron_model_runner.py:2071-2076` |
|
||
| `NCC_INLA001` collective reads an I/O tensor | copy into an intermediate buffer before the collective | `reduce_scatter_v.py:89,177,324` |
|
||
| Broadcast boolean AND miscompiles (`[4096,5120] & [1,5120]`) | fold validity into positions and use one `where` (CHRS-985) | `attention/attention_segmented_cte.py:546` |
|
||
| `NCC_ITEN406` "too many partition dims" (strided AP), often after graph breaks | keep fullgraph; avoid strided views | `docs/model-dev/debugging.md:150-160` |
|
||
| `NCC_ILSA901` transpose AP in fused MLP at T=2/4 | in-tree MLP path | `functional/qwen3_intree_mlp.py:174-202` |
|
||
| `NCC_INKI016` query seqlen must be a multiple of 128 | bucket validation | `utils/bucket_utils.py:455` |
|
||
| `NCC_IBTN243` SendRecv does not lower at lnc1 | prefill at lnc2, separate engine | `gpt-oss-gk-decode.md:19` |
|
||
| `-1` slot → OOB DGE scatter wedges the queue | `paged_slot_indices` clamps to null block | `model/kv_cache.py:80` |
|
||
| Hidden element offsets > 2**32 into giant buffers | allocate per-layer KV | `neuron_model_runner.py:10988` |
|
||
|
||
No `InferNonlocalTensors` hit was found in this repo.
|
||
|
||
## 7. FP8
|
||
|
||
- KV fp8 uses `torch.float8_e4m3fn` (`utils/dtype_utils.py:7-14`). Values are clamped to `FP8_CLAMP_MAX`: **240 on trn2** (the hardware/NKI `float8_e4m3`) and 448 on trn3 (`:16-41`). The constant is resolved at import so it is a trace-time constant. Write path: `(k*k_scale).clamp(±240).to(e4m3fn)` (`model/qwen3/model.py:311-321`).
|
||
- On trn2 the compiler flag `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` (inside the hlo2tensorizer string) makes the compiler treat e4m3fn as e4m3 (`neuron_model_runner_v2.py:673-686`).
|
||
- Static-fp8 weights (ModelOpt e4m3fn, range 448) are **downscaled by 240/448 at load** and the dequant scale is multiplied by 448/240 (`model/llama3/weight_loaders_static_fp8.py:45-80`, `model_static_fp8.py:126-160`). The prefix match `trn2` / `trn2-*` matters (`:150`).
|
||
- Some kernels and the swa path use 448 or 240 differently (`attention/swa_fused.py:45`). Check each kernel's range.
|
||
|
||
## 8. Patches relevant to torch/dynamo
|
||
|
||
- `vllm/patches/*` patches vLLM, not Dynamo or AOT (DP balancer, shm spin, pin_memory off, `accelerator.empty_cache` guard, port hold, replica mesh).
|
||
- Settings worth copying:
|
||
- `cache_size_limit=2**62` plus a `fail_on_recompile` stance at serve time.
|
||
- Warmup and serve must agree on grad mode, because Dynamo guards GLOBAL_STATE. Lite traces under `no_grad`, not `inference_mode` (`utils/neuron_utils.py:35-55`).
|
||
- Avoid logging inside traced code. Even `is_native_backend()` can log and break the graph, so it is read at construction (`nn/embedding.py:98-100`).
|
||
- Use precompile and AOTAutograd caching via `TORCH_CACHING_PRECOMPILE=1` and `TORCHINDUCTOR_AUTOGRAD_CACHE=1` (`vllm/inductor_cache.py:1-22`).
|
||
|
||
## 9. DeepSeek-V4 / V4.1 in this checkout
|
||
|
||
None. `contrib/models/` contains only `template/`, and `vllm_neuron/model/` has gpt_oss, llama3, qwen3, qwen3_vl, qwen3_asr and synthetic. "deepseek" appears only in comments (the V3.2 indexer key-only spec in `model/kv_cache.py:54`). PR #71 is unmerged here.
|
||
|
||
## What DeepInfer should copy (ranked by impact)
|
||
|
||
1. **KV as per-layer, non-overlapping, memoized views.** Allocate one tensor per layer (`(2,N,H,B,D)` or separate K and V). Bind them as module attributes. Never re-slice the same region twice. This removes the synthetic base and the `as_strided` path, and avoids the 2**32 offset bug.
|
||
2. **KV writes as plain `index_put_` with clamped indices** (`paged_slot_indices`, invalid slots go to null block 0). Rely on the backend's mutation-to-output aliasing (donation), not a custom op.
|
||
3. **One full-graph forward per (phase, bucket)**, including sampling and the in-graph state commit. Let neuronx-cc modular flow dedup layers (O1/O2), and try `-O3` for decode. Do not split per layer or per expert.
|
||
4. **Use the compiler flag set verbatim**: `--auto-cast=none`, a single `--internal-hlo2tensorizer-options` string (`--modular-flow-mac-threshold=10`, plus the fp8 flag on trn2), `--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop`, and `--model-type=generic` if `transformer` hurts.
|
||
5. **Host-built, fixed-shape metadata, H2D outside the graph.** Decode input_ids come from the on-device previous sample. No `.item()` in the graph. Use `fail_on_recompile` at serve time.
|
||
6. **Functional collectives with string reduceOp, SPMD rank tensor, no collective on graph I/O.** On trn2, verify all_to_all; vllm-neuron emulates it with all_gather.
|
||
7. **FP8: e4m3fn storage, clamp ±240, unsafe-e4m3fn-as-e4m3 flag**, and a 240/448 weight downscale for OCP checkpoints.
|
||
|
||
<details><summary>中文</summary>
|
||
|
||
# vllm-neuron 如何编译 decode 一步（给 DeepInfer Neuron 移植的笔记）
|
||
|
||
来源：本地 `Trn2/vllm-neuron`（HEAD `f32079d`，"Merge pull request #63 add-contrib-folder"）。路径相对仓库根目录。本文只读调研，没有改动 vllm-neuron 的任何内容。
|
||
|
||
**前提说明。** vllm-neuron **不用** `torch-neuronx`，也不用它的 `"neuron"` 后端。它通过 `libtorch-neuronx-lite`（`requirements/core.txt:7`）编译。lite 自带一份原生 lowering：Dynamo → AOTAutograd（带 Neuron decomps）→ torch-mlir → StableHLO → neuronx-cc（`docs/design/compilation/how_vllm_is_integrated_with_torchneuron.md`）。lite 源码不在这个仓库里。下文凡是讲 lite 内部行为的，依据都是注释和文档，不是代码。
|
||
|
||
## 1. 编译后端与编译参数
|
||
|
||
- `envs.get_compile_backend_name()`（`vllm_neuron/envs.py:789-803`）有两条路线：
|
||
- **原生路线**（默认；`VLLM_NEURON_BACKEND` 不设，或设为 `neuron_native`/`neuron`）：调用 `libtorch_neuronx_lite.compile.native_backend.register()`，返回它注册的名字。这个名字刻意**不叫** `"neuron"`，`"neuron"` 留给 torch_neuronx 包。
|
||
- **XLA 路线**（已弃用）和 CPU 模式：返回 `"neuron_libtorch"`，在 import 时注册（`vllm_neuron/__init__.py:~128`）。
|
||
- `torch-neuronx` 和 lite 要共存，必须设 `TORCH_DEVICE_BACKEND_AUTOLOAD=0`，否则报 "multiple backend fallbacks for PrivateUse1"。
|
||
- torch.compile 调用在 `vllm_neuron/vllm/worker/neuron_model_runner_v2.py:1903-1908`：
|
||
`torch.compile(NeuronForwardAdapter(model), backend=..., fullgraph=not DEBUG, options=self.compile_options)`。
|
||
- `torch._dynamo.config.cache_size_limit = 2**62`（`:1860`）。bucket 本来就要特化，重编译由 stance 来拦。
|
||
- serve 时的 forward 包在 `torch.compiler.set_stance("fail_on_recompile")` 里（`utils/neuron_utils.py:26-32`，在 `_serve_forward_guard` 用到，`neuron_model_runner_v2.py:1223`）。
|
||
- warm start 可选 `guard_filter_fn`（`envs.py:814`）：`keep_portable_guards_unsafe` 或 `skip_all_guards_unsafe`。
|
||
- `compile_options`（`neuron_model_runner_v2.py:387`、`:672-701`）：
|
||
- `alias_meta_to_neuron=True`：lite 把被改写的输入接回成别名输出。
|
||
- `compiler_args` 依次是：
|
||
- `--auto-cast=none`
|
||
- `--verbose=35`
|
||
- `-O{vLLM optimization_level}`
|
||
- `--internal-hlo2tensorizer-options=--modular-flow-mac-threshold=10`，trn2 上有 fp8 时在同一字符串里追加 ` --experimental-unsafe-fp8e4m3fn-as-fp8e4m3`
|
||
- `--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop`
|
||
- LNC 与编译器默认值不同时，再加显式 `--lnc=N`
|
||
- 需要 `--enable-nested-dynamic-loop`，是因为 NKI kernel 内联进父 BIR 后有嵌套 `scf.for`，birverifier 会拒多出来的回边（`neuron_model_runner.py:2071-2076`）。
|
||
- hlo2tensorizer 选项必须写在**同一个**字符串里。第二个 `--internal-hlo2tensorizer-options` 会覆盖第一个（`neuron_model_runner.py:2037-2047`）。
|
||
- V1 runner 给 GPT-OSS giga-kernel 图加了 `--enable-annotate-fsdp-cc=false` 和 `--model-type=generic`（`:2047`、`:2079`）。torch-neuronx 会注入 `--model-type transformer`，而 neuronx-cc 以最后出现的那个为准。
|
||
- **按图设参数。** 模型可以实现 `compiler_args_for(graph_type, example_inputs, base)`（`neuron_model_runner_v2.py:703-723`；`docs/design/framework/per_bucket_compiler_args.md`）。Qwen3-VL/ASR 的 decode 和 encoder 图用 **-O3**，理由是："-O3 关掉 neuronx-cc modular flow；-O1 默认开着它，会给每个 collective 包一层 PSEUDO_FUNCTION_RETURN/EVENT_SEMAPHORE"（`model/qwen3_vl/model_common.py:290-301`）。prefill 保持 -O2。
|
||
- 原生路线默认 O2，XLA 路线默认 O1（`docs/guides/how-to-choose-compilation-route.md`）。
|
||
|
||
## 2. KV cache 作为图输入（synthetic base 问题）
|
||
|
||
- **KV cache 不是 forward 参数。** `bind_kv_cache` 把普通 tensor 挂成每个 attention 模块的属性：`layer.self_attn.k_cache = kv[name][0]`、`.v_cache = kv[name][1]`（`model/qwen3/model.py:1160-1166`）。Dynamo 把这些属性提升为图输入。adapter 签名（`worker/neuron_forward_adapter.py:209-227`）里只有元数据和状态张量。
|
||
- **分配方式**（V2，`neuron_model_runner_v2.py:917-1215`）：
|
||
- 先建一个 `int8` backing（`:1140`）。
|
||
- 每层取 `raw = backing[start:start+span]`，再 `raw.view(dtype).view(2, num_blocks, kv_heads, block_size, head_size)`。**K = typed[0]，V = typed[1]**，是同一块每层区域上两个不重叠的 view。
|
||
- **避开 synthetic base 的关键。** view 对象做了记忆化：`shared_dtype_view` 和 `shared_kv_slices` 以 `(id(raw), dtype, kv_shape)` 为 key（`:1148-1186`）。共享字节的层因此拿到**完全相同的 tensor 对象**，Dynamo 会把它们去重成一个 placeholder。`:1160-1166` 的注释原话："对共享 storage 每层重新切片，会给 Dynamo 一组重叠输入，AOTAutograd 会把它们合并成整块 synthetic base，把 slot 大小的 KV 更新放大成整块写。"
|
||
- 也就是说，AOTAutograd 看到的 view 之间从不部分重叠：K 和 V 不相交，层与层不相交。AOTAutograd 的重叠检查会把它们当成各自独立的被改写输入，不会造 synthetic base，也不会回放 `as_strided`。
|
||
- host 端预检查会拒绝 `(dtype, heads, block, head_size)` 不一致的共享层，因为那样会产生不同的、相互重叠的 view 对象（`:1086-1110`）。
|
||
- **2**32 偏移隐患。** V1 runner 故意**每层单独 `torch.zeros`**，不从一个大 buffer 里切。原因是："指向多 GiB 共享 storage 的 view，元素偏移会超过 2**32，lowering 可能截断它（静默的跨层 KV 损坏）"（`neuron_model_runner.py:10988-10995`）。V2 是从单个 backing 里切的。**DeepInfer 用每层独立分配最稳。**
|
||
- **KV 写入就是图内普通的 `index_put_`**，不是自定义 op，也不是 kernel 内部更新 cache：
|
||
- `paged_slot_indices(slot_mapping, block_size, num_blocks)`（`model/kv_cache.py:80`）在做 `//`、`%` **之前**，把非法或 `-1` 的 slot 钳到 null block 0。`-1` 会变成越界的 DGE scatter，卡死 NEFF 执行队列。
|
||
- 然后 `k_cache.index_put_((block_idx, head_idx, pos_idx), k_flat)`（prefill 在 `model/qwen3/model.py:336-343`，decode 在 `:582-589`）。
|
||
- decode 的 NKI kernel 以 `update_cache=False` 调用，返回 `K_new`、`V_new`，再由模型自己 scatter 回去（`:525-589`，`qwen3_vl/utils/decode_kv.py:60` 的 `scatter_new_kv`）。
|
||
- **改写如何变成别名。**
|
||
- lite 的 `alias_meta_to_neuron=True` 把每个被改写的输入变成一个额外输出加一条 `input_output_alias`（donation）。XLA 路线上的对应做法见 `docs/design/compilation/aliasing_output_rewrite_pass.md`（被改写的输入追加进输出，得到 `io_map{out_idx: in_idx}`）和 `inplace_to_outofplace_pass.md`。NKI kernel 用 `operand_output_aliases` 声明同样的别名关系。
|
||
- vllm-neuron **没有用** `mark_static_address`、`keep_inference_input_mutations`，也没有自定义 `torch._functorch` 配置。
|
||
- **runner 状态也是同一套。** `token_ids.index_put_(...)` 在图内执行（`neuron_forward_adapter.py:190`），并返回别名 bundle。
|
||
- 下一步 decode 读的是采样结果的 **clone**，一个非别名输出（`:345-348`）。这样各 TP rank 上的采样 collective 保持有序。
|
||
- 输出 future 要一直持有到下一步（`neuron_model_runner_v2.py:1453-1460`），否则在飞的 NRT 输入切片会被释放，TP 会卡死。
|
||
|
||
## 3. 图粒度
|
||
|
||
- **一个图就是整个 forward**，按 (phase, bucket) 各一个。图里包括 input_ids 构造、所有 decoder 层（Python 循环被 Dynamo 展开，`model/qwen3/model.py:896`）、LM head、设备端采样和状态提交（`neuron_forward_adapter.py:272-356`）。
|
||
- prefill 和 decode 是两个图，以 `linear_idx` 是否为 None 区分。
|
||
- NEFF 数 = `|num_seqs_buckets| × (|decode_ctx_buckets|+1)`（`docs/design/vllm/decode-context-length-bucketing.md:111-116`）。
|
||
- **没有按层编译的区域**，也没有"一层编译一次、所有层复用"的机制。重复层的复用交给 **neuronx-cc modular flow**（O1/O2 下的 `--modular-flow-mac-threshold=10`）。对延迟敏感的 decode 图用 `-O3` 把它**关掉**。
|
||
- giga-kernel 路线走得更远：GPT-OSS 全部 36 层用一个 NKI kernel 跑完，单个 decode NEFF，权重在 `[num_layers, ...]` 轴上堆叠（`docs/design/vllm/gpt-oss-gk-decode.md`）。
|
||
- DeepSeek-V4.1 PR #71 约 120 个 NEFF 的做法（按 expert 分图、MoE 用 `@torch._dynamo.disable`、每层一次 host 同步）与仓库其它部分正好相反。见 `vllm_neuron_pr71_analysis.md`。
|
||
- 编译耗时的应对办法是：TP 组内并行 trace 和编译（`parallel_compile`，`neuron_model_runner_v2.py:2035-2056`）、编译缓存、SPMD 单 NEFF 校验。不靠拆图。
|
||
|
||
## 4. bucket 与静态形状
|
||
|
||
- 默认值（`utils/bucket_utils.py:36-126`）：
|
||
- prefill token bucket：从 128 起取 2 的幂，直到 `max_num_batched_tokens`。
|
||
- decode batch bucket：从 1 起取 2 的幂，直到 `max_num_seqs`。
|
||
- 可选：KV segment bucket（前缀缓存用）和 decode 上下文长度 bucket。
|
||
- `_select_bucket` 选能装下的最小 bucket（`neuron_model_runner_v2.py:903`）。
|
||
- padding 规则：
|
||
- 填充 token 用 `PAD_TOKEN=0`。填充 slot 用 `-1`，由 `paged_slot_indices` 重定向到 null block 0。
|
||
- block table 在每行有效块之后填 `-1` 哨兵，attention kernel 跳过越界部分。
|
||
- 填充的 decode 行 `row_to_slot=-1`（`worker/neuron_input_build.py:1-13`；`docs/design/vllm/padding-batching-block-kv.md`）。
|
||
- 元数据只在 host 上用 NumPy 构造，从不读设备状态。每一步在**图外**用 `torch.from_numpy(...).to(device)` 拷到设备（`neuron_model_runner_v2.py:2687-2712`）。
|
||
- Python int（`max_query_len`、`block_size`、`kv_segment_size`）有意放进元数据 dict，让它们参与图特化。
|
||
- decode 的 `input_ids` 在**图内**由上一步设备端采样结果构造，稳态 decode 没有 D2H。
|
||
- **图内没有 `.item()`。** host 只在取输出时读一次（对采样 future 做 `.cpu().tolist()`）。
|
||
- 标度另存一份 Python float，跟张量放在一起，因为 `tensor.item()` 会触发 graph break（`model/llama3/model.py:443`）。
|
||
|
||
## 5. 图内 collective
|
||
|
||
- 世界组和模型并行组都是 **Gloo**。collective 被 trace 进 NEFF，运行时从不单独派发（`envs.py:858-866`）。
|
||
- 用到的 API：
|
||
- `torch.distributed._functional_collectives.all_reduce(x, reduceOp="sum", group=...)`。reduceOp 必须是**字符串**，`ReduceOp` 对象会让 precompile 的 pickle 失败（`nn/embedding.py:256-260`、`nn/rpl.py:147-150`）。
|
||
- `all_gather_tensor`。
|
||
- `torch.ops._c10d_functional.all_gather_into_tensor`、`all_to_all_single`、`wait_tensor`（`functional/collectives/all_to_all.py:127-193`）。
|
||
- 也用了普通的就地 `dist.all_reduce`（经 vLLM `GroupCoordinator`），Dynamo 会把它改写成函数式版本。
|
||
- 限制：
|
||
- **trn2 上 all_to_all 是模拟的**：all_gather 加本地切片，理由是 "ENC 缺 Mesh 算法"（`all_to_all.py:10-33`）。
|
||
- 所有 rank 必须 trace 出逐字节相同的 HLO（SPMD）。用运行时 `rank` 张量，不用 Python int；依赖 rank 的切片改成 `index_select(arange+offset)`。每个组注册完整的 replica-group mesh（`docs/design/compilation/spmd_compilation.md`、`patches/replica_mesh_patch.py`、`parallel/neuron_group_coordinator.py:122-130`）。
|
||
- **collective 不能直接读 NEFF I/O 张量**（`NCC_INLA001`）。要先拷到中间 buffer（`functional/collectives/reduce_scatter_v.py:80-95, 318-330`）。
|
||
|
||
## 6. 已知 neuronx-cc 绕法
|
||
|
||
| 现象 | 绕法 | 位置 |
|
||
|---|---|---|
|
||
| 内联 NKI 后 birverifier 拒嵌套循环 | `--enable-verifier=false --enable-nested-dynamic-loop` | `neuron_model_runner.py:2071-2076` |
|
||
| `NCC_INLA001` collective 读了 I/O 张量 | collective 前先拷到中间 buffer | `reduce_scatter_v.py:89,177,324` |
|
||
| 广播布尔 AND 编错（`[4096,5120] & [1,5120]`） | 把有效性折进 position，只用一次 `where`（CHRS-985） | `attention/attention_segmented_cte.py:546` |
|
||
| `NCC_ITEN406` "too many partition dims"（跨步访问模式），常见于 graph break 之后 | 保持 fullgraph，避免跨步 view | `docs/model-dev/debugging.md:150-160` |
|
||
| `NCC_ILSA901` 融合 MLP 在 T=2/4 时 transpose 访问模式出错 | 走 in-tree MLP 路径 | `functional/qwen3_intree_mlp.py:174-202` |
|
||
| `NCC_INKI016` query seqlen 必须是 128 的倍数 | bucket 校验 | `utils/bucket_utils.py:455` |
|
||
| `NCC_IBTN243` lnc1 下 SendRecv 降不下去 | prefill 用 lnc2，单独起引擎 | `gpt-oss-gk-decode.md:19` |
|
||
| `-1` slot 变成越界 DGE scatter，卡死队列 | `paged_slot_indices` 钳到 null block | `model/kv_cache.py:80` |
|
||
| 大 buffer 里元素偏移超过 2**32 | KV 每层独立分配 | `neuron_model_runner.py:10988` |
|
||
|
||
本仓库里没有搜到 `InferNonlocalTensors`。
|
||
|
||
## 7. FP8
|
||
|
||
- KV fp8 用 `torch.float8_e4m3fn` 存储（`utils/dtype_utils.py:7-14`）。数值钳到 `FP8_CLAMP_MAX`：**trn2 上是 240**（硬件/NKI 的 `float8_e4m3`），trn3 上是 448（`:16-41`）。这个常量在 import 时就解析好，所以在 trace 时是常量。写入路径：`(k*k_scale).clamp(±240).to(e4m3fn)`（`model/qwen3/model.py:311-321`）。
|
||
- trn2 上的编译参数 `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`（放在 hlo2tensorizer 字符串里）让编译器把 e4m3fn 当 e4m3 处理（`neuron_model_runner_v2.py:673-686`）。
|
||
- 静态 fp8 权重（ModelOpt e4m3fn，范围 448）在**加载时乘 240/448 缩小**，反量化标度相应乘 448/240（`model/llama3/weight_loaders_static_fp8.py:45-80`、`model_static_fp8.py:126-160`）。平台判断要前缀匹配 `trn2` / `trn2-*`（`:150`）。
|
||
- 个别 kernel 和 swa 路径对 448/240 的用法不一样（`attention/swa_fused.py:45`），要逐个核对 kernel 的数值范围。
|
||
|
||
## 8. 与 torch/dynamo 相关的 patch
|
||
|
||
- `vllm/patches/*` 改的是 vLLM，不是 Dynamo/AOT（DP 负载均衡、shm 自旋、关掉 pin_memory、`accelerator.empty_cache` 保护、端口占位、replica mesh）。
|
||
- 值得照搬的设置：
|
||
- `cache_size_limit=2**62`，serve 时加 `fail_on_recompile` stance。
|
||
- warmup 和 serve 的 grad 模式必须一致，因为 Dynamo 会对 GLOBAL_STATE 加 guard。lite 在 `no_grad` 下 trace，不是 `inference_mode`（`utils/neuron_utils.py:35-55`）。
|
||
- traced 代码里不要打日志。连 `is_native_backend()` 都可能打日志并导致 graph break，所以它在构造时就读好（`nn/embedding.py:98-100`）。
|
||
- 用 `TORCH_CACHING_PRECOMPILE=1` 和 `TORCHINDUCTOR_AUTOGRAD_CACHE=1` 开 precompile 和 AOTAutograd 缓存（`vllm/inductor_cache.py:1-22`）。
|
||
|
||
## 9. 本仓库里的 DeepSeek-V4/V4.1
|
||
|
||
没有。`contrib/models/` 下只有 `template/`；`vllm_neuron/model/` 下是 gpt_oss、llama3、qwen3、qwen3_vl、qwen3_asr、synthetic。"deepseek" 只出现在注释里（V3.2 indexer 的 key-only spec，`model/kv_cache.py:54`）。PR #71 在这里没有合并。
|
||
|
||
## DeepInfer 该照搬什么（按影响排序）
|
||
|
||
1. **KV 按层独立、互不重叠、view 记忆化。** 每层一个张量（`(2,N,H,B,D)`，或 K、V 分开），挂成模块属性。同一块区域绝不切第二次。这样就没有 synthetic base 和 `as_strided` 路径，也躲开了 2**32 偏移的 bug。
|
||
2. **KV 写入用普通 `index_put_`，索引先钳位**（`paged_slot_indices`，非法 slot 指向 null block 0）。依赖后端把改写转成输出别名（donation），不写自定义 op。
|
||
3. **每个 (phase, bucket) 一个 full-graph forward**，采样和状态提交都放在图内。重复层交给 neuronx-cc modular flow 去重（O1/O2），decode 可以试 `-O3`。不要按层或按 expert 拆图。
|
||
4. **编译参数原样照抄**：`--auto-cast=none`；hlo2tensorizer 选项只写一个字符串（`--modular-flow-mac-threshold=10`，trn2 上再加 fp8 参数）；`--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop`；如果 `transformer` 有负面影响，加 `--model-type=generic`。
|
||
5. **元数据在 host 构造、形状固定、在图外做 H2D。** decode 的 input_ids 来自设备端上一步采样结果。图内不用 `.item()`。serve 时开 `fail_on_recompile`。
|
||
6. **函数式 collective + 字符串 reduceOp + SPMD rank 张量；collective 不直接读图 I/O。** trn2 上要实测 all_to_all，vllm-neuron 是用 all_gather 模拟的。
|
||
7. **FP8：e4m3fn 存储，钳 ±240，加 unsafe-e4m3fn-as-e4m3 参数**；OCP checkpoint 的权重按 240/448 缩小。
|
||
|
||
</details>
|