DeepInfer/docs/neuron/vllm_neuron_compile_notes.md

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).
• 232 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.

中文

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）。
• 232 偏移隐患。** 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 缩小。