DeepInfer/docs/neuron/vllm_neuron_pr71_analysis.md

# vllm-neuron PR #71 (DeepSeek-V4.1-Flash + DSpark on trn2): analysis for the DeepInfer Neuron port

Source: [vllm-project/vllm-neuron#71](https://github.com/vllm-project/vllm-neuron/pull/71), head `30d54d1` (OswaldHe/vllm-neuron), open, not merged, read 2026-10-22.
Paths below are relative to `contrib/models/deepseek_v41/src/` unless prefixed. `model.py:N` = line N at that SHA.
Note: the call-site line numbers in `kernels/__init__.py` are stale; the ones here come from `model.py` itself.

## 0. Summary

- **No XLA in the model code.** It never imports `torch_xla` or NxDI. It uses `libtorch_neuronx_lite` (`wrap_nki` for NKI, `model.py:401`), `torch.compile(backend=vllm_neuron.envs.get_compile_backend_name())` (`model.py:445`, `:2234`), `torch.ops._c10d_functional` collectives (`model.py:225-240`) and `vllm_neuron.*` helpers. Caveat: the numbers were measured on vllm-neuron 0.24.0.1.1.0, and the README lists `torch`/`torch-xla` 2.11 in that stack. In 0.24 the default compile route is `neuron_libtorch` (lite on top of XLA); the native route needs `VLLM_NEURON_BACKEND=neuron_native`. The PR never says which route it ran. Its error strings ("Detected non-contiguous slicing for requested Device Tensor") look like the native device. Either way, the kernels and the model logic do not depend on XLA.
- **Not a whole-model static graph.** About 120 NEFFs (`__init__.py:83`), one graph execution costs about 1.12 ms (`model.py:2353`), and every MoE layer syncs with the host (`model.py:2775`). The decode step (ITL) is about 320-370 ms. TPOT is about 117 ms only because DSpark accepts about 2.75 tokens per step (`TPOT × (1 + 5 × accept)` = step time, `model.py:283`).
- **Parallel layout.** Attention runs TP=8 and is replicated 8x. Experts are partitioned by **expert id** (6 of 384 per rank), so this is EP-style, not TP-sharded experts. Tokens are replicated, and there is one world all-reduce per MoE layer. Engram is row-partitioned across 64 ranks and resident in HBM. Everything lives in HBM: about 13.9 GiB per rank.
- **Kernels.** There are 12 NKI kernels, all written against `nki` 0.6 (`import nki`, `nki.isa` keyword `dst=` form, `@nki.jit`), the same `nki>=0.6` that our `requirements-neuron.txt` pins. There is no MoE kernel. Expert GEMMs are `blockquant.fp4_linear_2core`, run once per expert from a host loop.

## 1. Parallelism and placement

| Part | Layout | Ref |
|---|---|---|
| MLA attention | `ATTENTION_TENSOR_PARALLEL_SIZE=8` (capped by `o_groups=8`), so the attention is replicated `64/8 = 8` times. Each rank holds 8 of 64 query heads. `wq_b` is column-split by head, `wo_a` is split by o-group (1 local group), `wo_b` is row-split. `wq_a`, `wkv`, norms and `attn_sink` (sharded with the heads) are replicated. One fp32 all-reduce of `[T,5120]` over the 8-rank attention group, which is 8 consecutive ranks on 2 neighbouring devices. | `config.py:30-43`, `model.py:4099-4124`, `:1917-1946` |
| Indexer | `index_n_heads=32` split by 8 (4 per rank). `wq_b` and `weights_proj` are split by head, `wk`/`k_norm` are replicated. The score is all-reduced over the attention group, then top-k with `index_topk=512`. Only `index_source_layer_ids` [2,8,14,20,24,28,32,36] run an indexer; the other layers reuse `shared.topk_idxs`. | `config.py:117-124`, `model.py:1471-1605`, `:1782-1861` |
| Compressor / compressed KV | Replicated weights (bf16). Only `kv_source_layer_ids` [2,8,14,20] own a compressed cache. Later layers with the same ratio reuse the rows read by their source layer; this saves memory only, with ±1.3% latency. | `model.py:1412-1468`, `:1640-1690` |
| Dense, shared expert, router | Replicated on all 64 ranks: shared expert in fp8, gate in bf16, mHC `fn/base/scale`, norms. | `model.py:4127-4140` |
| Embedding / LM head | Vocab-parallel over 64 ranks (aligned to 128, padded). Embedding uses mask + world all-reduce; logits use a world all-gather. | `model.py:4468-4512` |
| Routed experts | Partitioned by id: rank r owns experts `[6r, 6r+6)` of 384 in each of the 40 backbone layers. The 3 MTP stages have 128 experts each (2 per rank). All are resident in HBM in fp4 (4.3 GiB per rank), byte-for-byte as in the checkpoint. Weight `[n, k/2]` uint8 holds two e2m1 values per byte; scale `[n, k/32]` uint8 holds ue8m0. Fp8 weights are `[n,k]` uint8 e4m3 with a `[ceil(n/32), ceil(k/32)]` ue8m0 scale. Gate and tokens are replicated, so every rank computes the same routing. It runs its own experts on its rows and scatter-adds into an fp32 carry, then **one world all-reduce** of `[T,5120]` fp32 runs. There is no dispatch collective. | `weight_loader.py:13-33`, `model.py:2616-2928`, `:657-706` |
| Engram | 2 tables (layers 1 and 14) of 384M rows × 256. The fp8 checkpoint is dequantized to **bf16 at load** and row-partitioned: about 6.0M rows per rank, 2.86 GiB per table, **in HBM**. Hashing runs on device (`NgramHashState`). Each rank looks up its own rows (`F.embedding` + mask), then a world all-reduce of `[1,T,24,256]` bf16 runs. A host streaming fallback exists but is not used. | `weight_loader.py:29-33`, `model.py:3112-3215`, `:4680-4706` |
| mHC | 4 streams. All coefficients and the 20-step Sinkhorn run locally (replicated). There are no collectives. | `model.py:3217-3283` |
| DSpark (`mtp.0-2`) | Same `Attention`/`MoE`/`BlockQuantLinear` classes, so the same layout. Markov head embed/head are vocab-parallel. | `model.py:3549-3740`, `:4141-4157` |

**Collectives per forward (batch 1):** each of the 40 layers has 1 attention-group all-reduce and 1 world all-reduce (MoE). The 8 indexer layers add 1 attention-group all-reduce each. The 2 engram layers add 1 world all-reduce each. Then there is 1 world all-reduce for the embedding and 1 world all-gather for the logits, and the DSpark stages add their own attention/MoE/embedding/markov collectives. This totals the "~43 world-wide all-reduces" that `model.py:248` cites. Measured cost: a 64-rank all-reduce takes about 1.9 ms against about 0.1 ms for 8 consecutive ranks. The cost is per-call latency, not bytes (`model.py:248-262`), so world collectives take about 80 ms per step. Factoring the world into 8×8 was 2.8% *slower*. TP×PP was probed (`DEEPSEEK_V41_PP_EMULATE`, `kernels/pp_permute.py`) but does not ship.

## 2. Execution model

- **Compile unit.** The runner compiles the model forward with graph breaks allowed (`VLLM_NEURON_DEBUG_MODE=1` gives `fullgraph=False`, `__init__.py:34`). The model also creates its own compiled entry points with `torch.compile(..., fullgraph=True, options=_COMPILE_OPTIONS)`:
  - one per NKI kernel (`_nki_entry`, `model.py:386-450`);
  - one per block for `attn_half`, which covers mHC, norm, attention, indexer, compressor and KV writes (`_block_entry`, `model.py:3432`, `:5025-5051`);
  - one per small step: router (`_gate_step`), each routed-expert step, shared expert, `_post_into`, and the collectives (`_compiled`, `model.py:2257`).

  `MoE.forward` is `@torch._dynamo.disable` (`model.py:2851`). Dynamo limits are raised to 2^30 because more than 256 variants otherwise make dynamo silently fall back to eager (`model.py:192-223`). Request-time compilation is allowed because the graph set depends on what the router picks (`model.py:5215-5249`).
- **Compiler flags** (`model.py:180-190`): `--auto-cast=none -O1`, `--modular-flow-mac-threshold=10`, `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` (required by the fp8 kernels), `--enable-verifier=false --enable-nested-dynamic-loop`.
- **Buckets.** There is no separate context-encoding / token-generation graph family; prefill and decode go through the same per-layer graphs, keyed by row count.
  - Prefill is segmented: `max_num_batched_tokens` is 4096 at 64 ranks (`__init__.py:451-561`). TTFT fits `ceil(tokens/segment) × F + tokens × ~1.06 ms`, with F ≈ 325-783 ms depending on segment and context.
  - Decode has `(1+5) × num_seqs` rows. `MAX_NUM_SEQS=1` because runs with more stalled (`config.py:45-49`).
  - Expert row lists are padded to powers of two with a floor of 128 (`model.py:2503`, `:2790`), because the GEMM tiles M at 128.
  - The compressed-KV read length is pinned to `max_model_len`, so prefill chunks and decode share graphs (`model.py:1178-1185`, `:4425-4440`). The cost is that a short request on a long-context engine reads 13x more.
- **KV layout.**
  - The window KV is vLLM's paged allocation, one per layer, `[blocks, 1, block_size, 512]` bf16. It is declared as full attention because the paged slots also index the other per-slot caches (`model.py:4268-4312`).
  - The model allocates flat per-slot buffers for kv-source layers: `compress_kv_cache [slots,512]` bf16, `index_k_cache [slots,128]` bf16, and compressor `pool_kv`/`pool_score [slots,512]` bf16 (`model.py:4314-4377`, `:1431-1439`).
  - Quantization is **emulated**: values pass through fp8 QDQ (window KV) or fp4 QDQ (compressed KV, index K/Q) and are **stored as bf16** (`model.py:1770`, `:1543-1557`, `:1823`).
  - Addressing goes through `PagedSpan` (`model.py:1302-1392`).
- **DSpark inside the target forward.** `forward` (`model.py:4515-4677`) runs 40 blocks and taps the attention inputs of layers 37/38/39 (`_dspark_tap`, mean over the 4 streams). It then runs the LM head, the on-device sampler and the rejection sampler. `_draft` → `forward_spec` (`model.py:5053-5145`) then runs the 3 MTP stages over a 5-token draft block (seed + noise tokens). Their attention reads the main window and the draft block's own KV, with no compressed KV. The last stage samples 5 tokens one at a time through a rank-256 Markov logit bias and adds a confidence head. Drafts are stored in `self.last_draft`, and `DSparkProposer.propose` (`drafter.py`) hands them to vLLM. There is no separate draft NEFF or warm-up. This needs the plugin's proposer registry (the `vllm_neuron/vllm/spec_decode` and runner hunks in this PR).
- **Host syncs are not avoided.**
  - Each MoE layer copies `indices.cpu()` to the host and builds row lists there (`model.py:2775-2800`).
  - Span and slot planning calls `.tolist()` on `slot_mapping` and `block_table` (`model.py:4766`, `:4831`, `:4943`).
  - `int(start_pos)` is read in `Block.forward` (`model.py:3332`).
  - The fixed cost per graph execution (about 1.12 ms) is the bottleneck the comments identify (`model.py:2351-2370`). Experts can be fused 2-3 per graph through `DEEPSEEK_V41_EXPERT_FUSE`, which is off by default.

## 3. Kernel inventory

Common to all kernels:
- `import nki; import nki.isa as nisa; import nki.language as nl`; entry points are `@nki.jit`.
- Outputs are allocated as `nl.shared_hbm` and **returned** (functional style). Scratch space is `nl.private_hbm`.
- 128-row tiles with a tail; reductions accumulate in fp32.
- Called through `wrap_nki(kernel)` (`libtorch_neuronx_lite.nki.nki_hop`) inside `torch.compile(fullgraph=True)`. `[2]` sets the grid for the 2-core variants.

"Ours" means the shim or op it replaces in DeepInfer (`vendor/` = `deepinfer/neuron_ops/vendor/`, `tl/` = `deepinfer/model_executor/layers/tilelang_ops/neuron/`, `v4/` = `deepinfer/model_executor/models/v4_modules/`).

| Kernel (entry, line) | Computes; I/O | Ours | Reuse on torch-neuronx 2.14 |
|---|---|---|---|
| `blockquant.fp8_linear` :399, `_raw` :417, `_2core` :477, `_raw_2core` :516 | Computes `x[m,k] @ W[n,k]^T`. The activation is first rounded through 1x32 fp8 (`_raw` skips this). The weight (uint8 e4m3 bytes plus a uint8 ue8m0 `[n/32,k/32]` scale) is decoded with bit arithmetic into **bf16 private HBM** on every call, then multiplied on the PE in bf16 (`nisa.nc_matmul`, fp32 PSUM). Output is `[ceil(m/128)*128, n]` bf16, and the caller slices it. `_2core` splits n across the 2 physical cores (`nl.program_id`, grid 2). | `vendor/deep_gemm/gemm.py:205` `fp8_fp4_gemm_nt` (dense, `v4/linear.py`), `einsum.py:77` `fp8_einsum` (`wo_a`, `v4/attn.py:215`), and the `tile_kernels.quant.per_token_cast` that comes before them | The kernel source can be reused as-is. Needs: the `nki_op` adapter; weights kept as checkpoint bytes (no `transform_sf_into_required_layout` repacking); the fp8 compiler flag; and grid=2 launch support on our stack for `_2core`. **Numerics: activation QDQ divides by 224, not 448** (Trn2 e4m3 saturates at 240), one mantissa bit off from DeepInfer's reference. |
| `blockquant.fp4_linear` :662, `fp4_linear_2core` :681 | Same pattern for fp4 experts: weight `[n,k/2]` holds two e2m1 values per byte, scale `[n,k/32]` ue8m0, output scaled by 2^64 after the matmul. 2core cut TPOT from 142.8 to 121.4 ms (`model.py:474-480`). | `gemm.py:316` `m_grouped_fp8_fp4_gemm_nt_contiguous` (`v4/mlp.py:119`) | Reusable. Called once per expert; our grouped API needs a loop or wrapper over local experts. |
| `blockquant.bf16_linear` :1221, `fp32w_linear` :1204 | Plain bf16 GEMM for the router, LM head, compressor, indexer `wk`/`weights_proj` and Markov head. `fp32w` narrows an **fp32 weight to bf16** in-kernel (mHC `fn`). | `gemm.py:107` `bf16_gemm_nt`; `vendor/deep_gemm/hc.py:10` `tf32_hc_prenorm_gemm` | Reusable. Precision note: bf16 instead of tf32 for the mHC prenorm GEMM. |
| `blockquant.*_fused`, `_tfused`, `_kmajor` | Variants that decode the weight in SBUF. `fp4_linear_fused` was measured **+50% TTFT / +24% TPOT** (`model.py:536-552`); the others are unreachable. | — | Skip. |
| `sparse_attn.gathered_attn_idx` :180 (live), `gathered_attn` :51 | For each token: indirect-DMA gather of K slots from `kv[n,512]` bf16 by `idx[T,K]` int32 (`kv.ap(vector_offset=…, indirect_dim=0)`); QK on bf16 operands into fp32; bias (column 0 = softmax scale, then -1e30 for empty slots); softmax with a learned sink; bf16 P·V. Input q is `[T*H,512]` bf16 with H=8 local heads; output `[T*H,512]` bf16. Requires K%128==0, d%128==0, H≤128. K = 128 window + 512 top-k = 640. | `vendor/flash_mla/flash_mla_interface.py:282` `flash_mla_with_kvcache`, `:155` `flash_mla_sparse_fwd` | The kernel can be reused, but **its KV layout differs**: it reads bf16 rows, while our cache stores fp8+ue8m0 records (528 B/token). Either switch the Neuron KV cache to bf16 (1024 B/token; fine at bs1 / 8K) or gather+dequantize in torch and call `gathered_attn`. Indirect DMA needs a device check under `nki_op`. |
| `index_score.weighted_scores` :41 | Computes Σ_h w[h,t]·relu(q_h·k). Inputs: q `[T*Hi,128]` bf16 (Hi=4 local), k `[n,128]` bf16 (n%128==0), w `[Hi,T]` fp32. Output `[T,n]` fp32 (PE matmul, then a second matmul for the head sum). | `vendor/deep_gemm/attention.py:123/225` `fp8_fp4_(paged_)mqa_logits` | Reusable once q/k are dequantized to bf16 (the PR runs fp4 QDQ in torch first). Top-k is **not** in the PR: it calls `vllm_neuron.functional.topk` (nkilib `gpsimd_topk` / vendored `rotational_topk`), which is also relevant to our #22 and `vendor/deep_select`. |
| `compress_pool.softmax_pool` :42 | Computes `(kv * softmax(score, dim=ratio)).sum(ratio)`; kv and score are `[G,ratio,512]` bf16, output `[G,512]` fp32. | `tl/compress.py:36/132` `compress_and_update_state_{prefill,decode}_tl` (pooling part) | Reusable. Small. |
| `hc_mixes.kernel` :66 | From raw `[R,24]` fp32 (the `fn` GEMM) and the stream `[R,4*5120]`: rms over the stream, then pre = σ+eps, post = 2σ, comb = softmax + **20 Sinkhorn iterations**. Outputs pre/post `[R,4]` and comb `[R,16]`, all fp32. | `tile_kernels.mhc` `norm_fn` / `pre_split_mixes` / `sinkhorn` (`v4/mhc_ops.py:10-16`) | Reusable. Check the iteration count and eps against our config (`sinkhorn_normalize_ref` defaults to 10 iterations). |
| `hc_expand.kernel` :28 | `[B,S,D]` → 4 copies `[B,S,4,D]` plus a one-hot pre_mix `[B,S,4]` fp32. DMA only. | `tile_kernels.mhc.expand_kernel.expand_to_mhc_fwd` | Reusable. |
| `hc_collapse.kernel` :24 | Σ_c pre[c]·x[c]: `[B,N,4,D]`, `[B,N,4]` fp32 → `[B,N,D]`. | `mhc_pre_apply_mix_fwd` | Reusable. |
| `hc_attn_out.kernel` :24 | post[j]·x + Σ_i comb[i,j]·res[i]: x `[B,S,D]`, res `[B,S,4,D]`, post `[B,S,4]`, comb `[B,S,4,4]` → `[B,S,4,D]`. | `mhc_post_fwd` | Reusable. |
| `rmsnorm.rms_norm` :30 | x `[R,D]`, weight `(1,D)`, eps as a `(1,1)` tensor → `[R,D]` in the input dtype. fp32 reduction plus one Newton step on rsqrt. | `tl/norm.py:75` `norm_tl`, `v4/norm.py` | Reusable. Check parity against `rstd_1ulp`. |
| `swiglu.swiglu` :40 | Clamped SwiGLU (up clamped to ±10, gate to ≤10) × routing weight: gate/up `[R,I]`, weights `(R,1)` or `(1,1)`, limit `(1,1)` → `[R,I]`. | `tile_kernels.quant.swiglu_forward_and_per_token_cast` (`v4/mlp.py:19`) | Reusable. It does not cast; the fp8 rounding is done by the next GEMM's QDQ. |
| `engram_gate.gated_add` :41 | For each of the 4 streams: rms(h)·rms(key) normalized dot → sign·sqrt → sigmoid gate; h + gate·value. h/key `[R,4*5120]`, value `[R,5120]`, weight `(1,4*5120)` = q_weight·k_weight. | `tile_kernels.engram.engram_gate_fwd` (`v4/engram.py:155`) | Reusable after checking it against our fused-weight semantics. |
| `pp_permute.kernel` :395 | `nki.collectives.collective_permute` plus `nisa.core_barrier` stage hand-off. **Not called**: the CC runtime refuses its replica group. | — | Skip. |

**MoE GEMM and dequant without a MoE kernel** (`model.py:2287-2380`, `:2851-2928`):
1. Router: `bf16_linear` + sqrt-softplus + nkilib top-k.
2. The host copies the indices back and builds a per-local-expert row list padded to a bucket of at least 128.
3. One compiled graph per expert: `index_select` rows → `fp4_linear_2core` (w1, w3) → `swiglu` → `fp4_linear_2core` (w2) → `index_add` into the fp32 carry.
4. World all-reduce.
5. Shared expert: the same chain with `fp8_linear_2core`, applied to all rows.

Dequantization happens inside every GEMM (weight bytes → bf16 in private HBM); there is no persistent bf16 copy.

## 4. Where the time goes and what makes TPOT ≈ 120 ms possible

The PR has no profile breakdown. The comments give these pieces:
- Graph execution: about 1.12 ms fixed each, about 120 per forward. A 128-row expert GEMM's roofline is about 154 µs, so most of the routed-expert time is overhead, not arithmetic (`model.py:2353-2363`).
- World all-reduces: about 43 × 1.9 ms ≈ 80 ms per step.
- Decode step ≈ 320-370 ms for 6 rows. TPOT = step / (1 + accepted), and the acceptance length is 2.75 on sonnet (36% per-position average).
- Prefill: F ≈ 0.3-0.8 s per forward plus about 1.06 ms/token.

The decisions that matter:
1. **Everything resident in HBM** (13.9 GiB per rank): no host traffic during the forward.
2. **Attention TP=8, replicated**, so the attention reduction stays within 2 devices.
3. **Experts by id with replicated tokens**: no dispatch, one all-reduce.
4. **Using both physical cores** of each LNC=2 core in the GEMMs: fp4 then fp8 2core brought TPOT from 142.8 to 116.9 ms.
5. **NKI for every op family**, to cut the number of graphs and ops.
6. **Prefill segment 4096.**
7. **DSpark** (≈2.75x).

Things they measured as no win: SBUF-fused dequant, the shared compressed read, `NEURON_RT_IO_RING_CACHE_SIZE`, 8×8 collective factoring, and `_EXPERT_ROW_FLOOR` below 128.

## 5. Integration plan for our port (ordered by expected payoff)

1. **#9 static graphs + #12 placement (structural, from about 150 s to under a second).**
   - Compile one decode graph per bucket with the NKI ops inside it.
   - Remove every host sync. To avoid PR #71's per-layer `.cpu()`, run **all 6 local experts over the bucketed rows**, with routing weight 0 for rows that are not routed to them. Their GEMM pads M to 128 anyway, so at decode this costs the same instructions per expert as their 128 bucket, and the graph stays static.
   - At bs1, consider **attention TP=8 groups instead of attention DP**: today one rank does all 64 heads, and the other 63 wait.
   - Once tokens are replicated, our all-gather dispatch is a no-op and the combine becomes one world all-reduce, like theirs. Measure all-reduce against reduce-scatter.
   - Copy their `_COMPILE_OPTIONS` flags. Budget about 1.9 ms per world collective.
2. **#10 GEMM: reuse `blockquant`.**
   - Put `fp8_linear(_raw)_2core` behind `fp8_fp4_gemm_nt`/`fp8_einsum`, `fp4_linear_2core` behind `m_grouped_fp8_fp4_gemm_nt_contiguous` (loop over local experts), and `bf16_linear`/`fp32w_linear` behind bf16/hc GEMMs.
   - Keep weights as checkpoint uint8 bytes (this also follows our fp8-as-bytes rule).
   - Decide whether to accept the 224 QDQ or patch the kernel to 448 and re-test the saturation flag.
   - Highest kernel payoff.
3. **#15 DSpark.** Their TPOT includes a 2.75x multiplier. Without it we need a step of 117 ms or less, not about 320 ms. The design to follow is in-forward drafting (taps from layers 37-39, a 5-token block, a Markov head).
4. **#11 attention/indexer.**
   - Choose the KV layout: bf16 cache plus `gathered_attn_idx`, or keep fp8 records with gather+dequant plus `gathered_attn`.
   - Put `index_score` behind `fp8_fp4_(paged_)mqa_logits`, and `compress_pool` behind `tl/compress.py`.
   - Borrow `vllm_neuron.functional.topk` (nkilib) for `deep_select.topk` and #22.
5. **#14 Engram.** Copy their placement: row-partitioned bf16 (or fp8, 2.9 GiB) tables in HBM, on-device hash and gather, and a world all-reduce. This removes the host gather and H2D entirely. It needs HBM headroom (#21): their dense part is 3.9 GiB per rank because attention is TP=8 and the vocab is split 64 ways. Add `engram_gate`.
6. **#13 small ops.** `hc_mixes`, `hc_expand`, `hc_collapse`, `hc_attn_out`, `rmsnorm` and `swiglu` are each small, but together they remove most of the per-layer op count. Do them after #9 so the profile is real.

**Things to check before reusing any kernel on our stack:**
- Whether `torch_neuronx.nki_op` (2.14.3) accepts `@nki.jit` kernels that *return* `shared_hbm` outputs, use `private_hbm` scratch, run under both eager and `torch.compile(backend="neuron")`, and launch with grid=2 at LNC=2.
- Whether indirect DMA (`.ap(vector_offset=…)`) works through `nki_op`.
- The shape asserts (k%128, K%128, heads≤128).
- Numeric parity against our CPU shims within quantization noise.

**License.**
- vllm-neuron is Apache-2.0 (`LICENSE`; `NOTICE`: "Copyright Amazon.com, Inc. or its affiliates"). DeepInfer is Apache-2.0 too, so the licenses are compatible.
- The kernel files have no per-file header. Only `kernels/__init__.py`, `model.py`, `weight_loader.py` and similar files carry `SPDX-License-Identifier: Apache-2.0`.
- When copying, add an SPDX header and an origin line to each copied file (repo, PR #71, SHA `30d54d1`, author Zifan He / OswaldHe), mark modified files as changed (Apache-2.0 §4b), and carry the relevant NOTICE text into our NOTICE or third-party notices.
- The PR is unmerged, so pin the SHA and re-check after merge.

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

# vllm-neuron PR #71（trn2 上的 DeepSeek-V4.1-Flash + DSpark）：面向 DeepInfer Neuron 移植的分析

来源：[vllm-project/vllm-neuron#71](https://github.com/vllm-project/vllm-neuron/pull/71)，head `30d54d1`（OswaldHe/vllm-neuron），状态 open、未合并，阅读日期 2026-10-22。
下文路径若无前缀，都相对于 `contrib/models/deepseek_v41/src/`。`model.py:N` 指该 SHA 下的第 N 行。
注意：`kernels/__init__.py` 里写的调用点行号已经过时，本文行号直接取自 `model.py`。

## 0. 摘要

- **模型代码里没有 XLA。** 不 import `torch_xla`，也不用 NxDI。它用的是 `libtorch_neuronx_lite`（NKI 走 `wrap_nki`，`model.py:401`）、`torch.compile(backend=vllm_neuron.envs.get_compile_backend_name())`（`model.py:445`、`:2234`）、`torch.ops._c10d_functional` collective（`model.py:225-240`），以及 `vllm_neuron.*` 工具函数。需要说明的是：性能数字测于 vllm-neuron 0.24.0.1.1.0，README 列出的这套环境里有 `torch`/`torch-xla` 2.11。0.24 的默认编译路线是 `neuron_libtorch`（lite 跑在 XLA 之上），要走 native 路线得设 `VLLM_NEURON_BACKEND=neuron_native`。PR 没写明用的是哪条路线；从报错字符串（"Detected non-contiguous slicing for requested Device Tensor"）看像是 native device。无论哪条路线，kernel 和模型逻辑都不依赖 XLA。
- **不是整模型静态图。** 约 120 个 NEFF（`__init__.py:83`），每次图执行约 1.12 ms（`model.py:2353`），每个 MoE 层都要和 host 同步一次（`model.py:2775`）。decode 一步（ITL）约 320-370 ms。TPOT 能到约 117 ms，只是因为 DSpark 每步平均接受约 2.75 个 token（`TPOT × (1 + 5 × accept)` = 一步时间，`model.py:283`）。
- **并行布局。** attention 用 TP=8，并复制 8 份。专家按**专家 id** 划分（每 rank 拿 384 个中的 6 个），也就是 EP 式，不是对专家做 TP 切分。token 在所有 rank 上复制，每个 MoE 层做一次全局 all-reduce。Engram 按行切到 64 个 rank，常驻 HBM。全部权重都在 HBM 里，每 rank 约 13.9 GiB。
- **kernel。** 共 12 个 NKI kernel，全部基于 `nki` 0.6 编写（`import nki`、`nki.isa` 的 `dst=` 关键字写法、`@nki.jit`），与我们 `requirements-neuron.txt` 锁定的 `nki>=0.6` 是同一个包。没有 MoE kernel。专家 GEMM 用的是 `blockquant.fp4_linear_2core`，由 host 循环逐个专家调用。

## 1. 并行与放置

| 部分 | 布局 | 引用 |
|---|---|---|
| MLA attention | `ATTENTION_TENSOR_PARALLEL_SIZE=8`（上限由 `o_groups=8` 决定），所以 attention 复制 `64/8 = 8` 份。每 rank 持有 64 个 query head 中的 8 个。`wq_b` 按 head 列切，`wo_a` 按 o-group 切（本地 1 组），`wo_b` 行切。`wq_a`、`wkv`、各 norm 和 `attn_sink`（随 head 一起切）都复制。每层在 8-rank attention 组上做一次 `[T,5120]` 的 fp32 all-reduce，这个组是相邻 2 个 device 上连续的 8 个 rank。 | `config.py:30-43`、`model.py:4099-4124`、`:1917-1946` |
| Indexer | `index_n_heads=32` 切 8 份（每 rank 4 个）。`wq_b` 和 `weights_proj` 按 head 切，`wk`/`k_norm` 复制。score 在 attention 组上 all-reduce，然后做 `index_topk=512` 的 top-k。只有 `index_source_layer_ids` [2,8,14,20,24,28,32,36] 这几层跑 indexer，其他层复用 `shared.topk_idxs`。 | `config.py:117-124`、`model.py:1471-1605`、`:1782-1861` |
| Compressor / 压缩 KV | 权重复制（bf16）。只有 `kv_source_layer_ids` [2,8,14,20] 拥有压缩 cache。后面同 ratio 的层复用源层读出的行；这只省内存，延迟变化在 ±1.3% 以内。 | `model.py:1412-1468`、`:1640-1690` |
| 稠密层、共享专家、router | 在全部 64 个 rank 上复制：共享专家为 fp8，gate 为 bf16，还有 mHC 的 `fn/base/scale` 和各 norm。 | `model.py:4127-4140` |
| Embedding / LM head | 词表在 64 个 rank 上切分（按 128 对齐并补齐）。embedding 用掩码加全局 all-reduce；logits 用全局 all-gather。 | `model.py:4468-4512` |
| 路由专家 | 按 id 划分：40 个 backbone 层中，每层 384 个专家里 rank r 拥有 `[6r, 6r+6)`。3 个 MTP stage 每个有 128 个专家（每 rank 2 个）。全部以 fp4 常驻 HBM（每 rank 4.3 GiB），字节与 checkpoint 完全一致。权重 `[n, k/2]` uint8，每字节存两个 e2m1；标度 `[n, k/32]` uint8，存 ue8m0。fp8 权重为 `[n,k]` uint8 e4m3，配 `[ceil(n/32), ceil(k/32)]` 的 ue8m0 标度。gate 和 token 都是复制的，所以每个 rank 算出的路由完全相同。各 rank 只对路由到自己专家的行做计算，scatter-add 到 fp32 累加器，然后做**一次全局 all-reduce**（`[T,5120]` fp32）。没有 dispatch collective。 | `weight_loader.py:13-33`、`model.py:2616-2928`、`:657-706` |
| Engram | 2 张表（第 1、14 层），每张 384M 行 × 256。fp8 的 checkpoint 在**加载时反量化为 bf16**并按行切分：每 rank 约 6.0M 行，每张表 2.86 GiB，**放在 HBM**。hash 在 device 上计算（`NgramHashState`）。各 rank 查自己那部分行（`F.embedding` 加掩码），然后对 `[1,T,24,256]` bf16 做全局 all-reduce。代码里有 host 流式回退路径，但实际没用。 | `weight_loader.py:29-33`、`model.py:3112-3215`、`:4680-4706` |
| mHC | 4 条流。所有系数和 20 步 Sinkhorn 都在本地计算（复制）。没有 collective。 | `model.py:3217-3283` |
| DSpark（`mtp.0-2`） | 用的是同样的 `Attention`/`MoE`/`BlockQuantLinear` 类，所以布局相同。Markov head 的 embed/head 按词表切分。 | `model.py:3549-3740`、`:4141-4157` |

**每次 forward 的 collective（batch 1）：** 40 层中每层有 1 次 attention 组 all-reduce 和 1 次全局 all-reduce（MoE）。8 个 indexer 层各多 1 次 attention 组 all-reduce。2 个 engram 层各多 1 次全局 all-reduce。另外还有 embedding 的 1 次全局 all-reduce、logits 的 1 次全局 all-gather，DSpark stage 也有自己的 attention/MoE/embedding/markov collective。合起来就是 `model.py:248` 说的"约 43 次全局 all-reduce"。实测成本：64-rank all-reduce 约 1.9 ms，连续 8 rank 约 0.1 ms。开销来自每次调用的固定延迟，与字节数无关（`model.py:248-262`），所以全局 collective 每步约占 80 ms。把全局拆成 8×8 反而**慢** 2.8%。TP×PP 做过探测（`DEEPSEEK_V41_PP_EMULATE`、`kernels/pp_permute.py`），但没有上线。

## 2. 执行模型

- **编译单元。** runner 编译模型 forward 时允许 graph break（`VLLM_NEURON_DEBUG_MODE=1` 对应 `fullgraph=False`，`__init__.py:34`）。模型自己还用 `torch.compile(..., fullgraph=True, options=_COMPILE_OPTIONS)` 建了一批编译入口：
  - 每个 NKI kernel 一个（`_nki_entry`，`model.py:386-450`）；
  - 每个 block 的 `attn_half` 一个，覆盖 mHC、norm、attention、indexer、compressor 和 KV 写入（`_block_entry`，`model.py:3432`、`:5025-5051`）；
  - 每个小步骤一个：router（`_gate_step`）、每次路由专家计算、共享专家、`_post_into`，以及各 collective（`_compiled`，`model.py:2257`）。

  `MoE.forward` 标了 `@torch._dynamo.disable`（`model.py:2851`）。dynamo 的上限调到了 2^30，因为图变体超过 256 个后，dynamo 会静默退回 eager（`model.py:192-223`）。由于图的集合取决于 router 的选择，所以允许在请求时编译（`model.py:5215-5249`）。
- **编译参数**（`model.py:180-190`）：`--auto-cast=none -O1`、`--modular-flow-mac-threshold=10`、`--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`（fp8 kernel 必需）、`--enable-verifier=false --enable-nested-dynamic-loop`。
- **bucket。** 没有单独的 context-encoding / token-generation 两套图；prefill 和 decode 走同一组按行数做 key 的逐层图。
  - prefill 分段：64 rank 时 `max_num_batched_tokens` 为 4096（`__init__.py:451-561`）。TTFT 符合 `ceil(tokens/segment) × F + tokens × ~1.06 ms`，F 随分段大小和上下文长度在 325-783 ms 之间。
  - decode 每步 `(1+5) × num_seqs` 行。`MAX_NUM_SEQS=1`，因为大于 1 时会卡住（`config.py:45-49`）。
  - 专家的行列表补齐到 2 的幂、下限 128（`model.py:2503`、`:2790`），因为 GEMM 在 M 方向按 128 分块。
  - 压缩 KV 的读取长度固定为 `max_model_len`，这样 prefill 分块和 decode 能共用图（`model.py:1178-1185`、`:4425-4440`）。代价是短请求跑在长上下文引擎上时要多读 13 倍。
- **KV 布局。**
  - 窗口 KV 用的是 vLLM 的分页分配，每层一份，`[blocks, 1, block_size, 512]` bf16。之所以声明成 full attention，是因为分页的 slot 同时用来索引其他逐 slot cache（`model.py:4268-4312`）。
  - 模型为 kv-source 层自己分配平铺的逐 slot buffer：`compress_kv_cache [slots,512]` bf16、`index_k_cache [slots,128]` bf16，以及 compressor 的 `pool_kv`/`pool_score [slots,512]` bf16（`model.py:4314-4377`、`:1431-1439`）。
  - 量化是**模拟**的：数值经过 fp8 QDQ（窗口 KV）或 fp4 QDQ（压缩 KV、index K/Q）后**以 bf16 存储**（`model.py:1770`、`:1543-1557`、`:1823`）。
  - 寻址通过 `PagedSpan` 完成（`model.py:1302-1392`）。
- **DSpark 在 target forward 内部。** `forward`（`model.py:4515-4677`）跑完 40 个 block，并抽取第 37/38/39 层的 attention 输入（`_dspark_tap`，对 4 条流取平均）。然后跑 LM head、device 上的 sampler 和拒绝采样。接着 `_draft` → `forward_spec`（`model.py:5053-5145`）在一个 5 token 的草稿块（种子 token 加噪声 token）上跑 3 个 MTP stage。这些 stage 的 attention 读主窗口和草稿块自己的 KV，不读压缩 KV。最后一个 stage 通过 rank-256 的 Markov logit 偏置逐个采样 5 个 token，并加一个 confidence head。草稿存在 `self.last_draft`，由 `DSparkProposer.propose`（`drafter.py`）交给 vLLM。没有单独的草稿 NEFF，也不用单独 warm-up。这需要插件的 proposer 注册机制（即本 PR 中 `vllm_neuron/vllm/spec_decode` 和 runner 的改动）。
- **没有消除 host 同步。**
  - 每个 MoE 层都把 `indices.cpu()` 拷回 host，在 host 上建行列表（`model.py:2775-2800`）。
  - span/slot 规划对 `slot_mapping` 和 `block_table` 调用 `.tolist()`（`model.py:4766`、`:4831`、`:4943`）。
  - `Block.forward` 里读了 `int(start_pos)`（`model.py:3332`）。
  - 注释指出的瓶颈是每次图执行的固定开销（约 1.12 ms）（`model.py:2351-2370`）。可以用 `DEEPSEEK_V41_EXPERT_FUSE` 把 2-3 个专家合进一张图，默认关闭。

## 3. Kernel 清单

所有 kernel 的共同点：
- `import nki; import nki.isa as nisa; import nki.language as nl`；入口为 `@nki.jit`。
- 输出用 `nl.shared_hbm` 分配并**直接返回**（函数式写法）。临时空间用 `nl.private_hbm`。
- 按 128 行分块并处理尾块；归约在 fp32 中累加。
- 通过 `wrap_nki(kernel)`（`libtorch_neuronx_lite.nki.nki_hop`）在 `torch.compile(fullgraph=True)` 内调用。2core 变体用 `[2]` 指定 grid。

"我们"一栏指它在 DeepInfer 里替换的 shim 或算子（`vendor/` = `deepinfer/neuron_ops/vendor/`，`tl/` = `deepinfer/model_executor/layers/tilelang_ops/neuron/`，`v4/` = `deepinfer/model_executor/models/v4_modules/`）。

| Kernel（入口、行号） | 计算内容；输入/输出 | 我们 | 在 torch-neuronx 2.14 上复用 |
|---|---|---|---|
| `blockquant.fp8_linear` :399、`_raw` :417、`_2core` :477、`_raw_2core` :516 | 计算 `x[m,k] @ W[n,k]^T`。激活先做 1x32 的 fp8 舍入（`_raw` 跳过这一步）。权重（uint8 e4m3 字节加 uint8 ue8m0 的 `[n/32,k/32]` 标度）每次调用都用位运算解码到 **bf16 private HBM**，再在 PE 上做 bf16 matmul（`nisa.nc_matmul`，fp32 PSUM）。输出 `[ceil(m/128)*128, n]` bf16，由调用方切片。`_2core` 把 n 切到 2 个物理核（`nl.program_id`，grid 2）。 | `vendor/deep_gemm/gemm.py:205` `fp8_fp4_gemm_nt`（稠密层，`v4/linear.py`）、`einsum.py:77` `fp8_einsum`（`wo_a`，`v4/attn.py:215`），以及它们之前的 `tile_kernels.quant.per_token_cast` | kernel 源码可原样复用。需要：`nki_op` 适配层；权重保持 checkpoint 字节布局（不做 `transform_sf_into_required_layout` 重排）；fp8 编译参数；`_2core` 还需要我们的栈支持 grid=2 发射。**数值：激活 QDQ 除以 224 而不是 448**（Trn2 的 e4m3 在 240 饱和），比 DeepInfer 参考实现少一位尾数。 |
| `blockquant.fp4_linear` :662、`fp4_linear_2core` :681 | fp4 专家用同样的模式：权重 `[n,k/2]` 每字节存两个 e2m1，标度 `[n,k/32]` ue8m0，matmul 之后乘 2^64。2core 把 TPOT 从 142.8 ms 降到 121.4 ms（`model.py:474-480`）。 | `gemm.py:316` `m_grouped_fp8_fp4_gemm_nt_contiguous`（`v4/mlp.py:119`） | 可复用。它每个专家调用一次；我们的 grouped API 需要在本地专家上循环或包一层。 |
| `blockquant.bf16_linear` :1221、`fp32w_linear` :1204 | 普通 bf16 GEMM，用于 router、LM head、compressor、indexer 的 `wk`/`weights_proj` 和 Markov head。`fp32w` 在 kernel 内把 **fp32 权重收窄为 bf16**（mHC 的 `fn`）。 | `gemm.py:107` `bf16_gemm_nt`；`vendor/deep_gemm/hc.py:10` `tf32_hc_prenorm_gemm` | 可复用。精度注意：mHC 的 prenorm GEMM 从 tf32 变成了 bf16。 |
| `blockquant.*_fused`、`_tfused`、`_kmajor` | 在 SBUF 里解码权重的变体。`fp4_linear_fused` 实测 **TTFT +50%、TPOT +24%**（`model.py:536-552`）；其余变体代码路径根本走不到。 | — | 跳过。 |
| `sparse_attn.gathered_attn_idx` :180（实际使用）、`gathered_attn` :51 | 逐 token：按 `idx[T,K]` int32 从 `kv[n,512]` bf16 中间接 DMA gather K 个 slot（`kv.ap(vector_offset=…, indirect_dim=0)`）；QK 用 bf16 操作数、fp32 累加；bias（第 0 列为 softmax scale，空 slot 为 -1e30）；带可学习 sink 的 softmax；P·V 用 bf16。输入 q 为 `[T*H,512]` bf16，H=8 个本地 head；输出 `[T*H,512]` bf16。要求 K%128==0、d%128==0、H≤128。K = 128 窗口 + 512 top-k = 640。 | `vendor/flash_mla/flash_mla_interface.py:282` `flash_mla_with_kvcache`、`:155` `flash_mla_sparse_fwd` | kernel 可复用，但**KV 布局不同**：它读 bf16 行，而我们的 cache 存的是 fp8+ue8m0 记录（每 token 528 B）。两种办法：把 Neuron 的 KV cache 改成 bf16（每 token 1024 B，bs1 / 8K 下没问题）；或者先在 torch 里 gather+反量化，再调 `gathered_attn`。间接 DMA 在 `nki_op` 下能否工作需要上机验证。 |
| `index_score.weighted_scores` :41 | 计算 Σ_h w[h,t]·relu(q_h·k)。输入：q `[T*Hi,128]` bf16（Hi=4 个本地 head）、k `[n,128]` bf16（n%128==0）、w `[Hi,T]` fp32。输出 `[T,n]` fp32（PE matmul，再用第二次 matmul 对 head 求和）。 | `vendor/deep_gemm/attention.py:123/225` `fp8_fp4_(paged_)mqa_logits` | 把 q/k 反量化成 bf16 后即可复用（PR 是先在 torch 里做 fp4 QDQ）。top-k **不在** PR 里：它调用 `vllm_neuron.functional.topk`（nkilib `gpsimd_topk` / vendored 的 `rotational_topk`），这对我们的 #22 和 `vendor/deep_select` 也有参考价值。 |
| `compress_pool.softmax_pool` :42 | 计算 `(kv * softmax(score, dim=ratio)).sum(ratio)`；kv 和 score 为 `[G,ratio,512]` bf16，输出 `[G,512]` fp32。 | `tl/compress.py:36/132` `compress_and_update_state_{prefill,decode}_tl`（其中的 pooling 部分） | 可复用。体量小。 |
| `hc_mixes.kernel` :66 | 输入 raw `[R,24]` fp32（`fn` GEMM 的结果）和整条流 `[R,4*5120]`：先对整条流求 rms，然后 pre = σ+eps、post = 2σ、comb = softmax 加 **20 次 Sinkhorn**。输出 pre/post `[R,4]` 和 comb `[R,16]`，全为 fp32。 | `tile_kernels.mhc` 的 `norm_fn` / `pre_split_mixes` / `sinkhorn`（`v4/mhc_ops.py:10-16`） | 可复用。需要对照我们的配置核对迭代次数和 eps（`sinkhorn_normalize_ref` 默认 10 次）。 |
| `hc_expand.kernel` :28 | `[B,S,D]` → 4 份拷贝 `[B,S,4,D]`，外加 one-hot 的 pre_mix `[B,S,4]` fp32。只做 DMA。 | `tile_kernels.mhc.expand_kernel.expand_to_mhc_fwd` | 可复用。 |
| `hc_collapse.kernel` :24 | Σ_c pre[c]·x[c]：`[B,N,4,D]`、`[B,N,4]` fp32 → `[B,N,D]`。 | `mhc_pre_apply_mix_fwd` | 可复用。 |
| `hc_attn_out.kernel` :24 | post[j]·x + Σ_i comb[i,j]·res[i]：x `[B,S,D]`、res `[B,S,4,D]`、post `[B,S,4]`、comb `[B,S,4,4]` → `[B,S,4,D]`。 | `mhc_post_fwd` | 可复用。 |
| `rmsnorm.rms_norm` :30 | x `[R,D]`、weight `(1,D)`、eps 为 `(1,1)` 张量 → `[R,D]`，dtype 与输入相同。fp32 归约，rsqrt 后再做一步 Newton 迭代。 | `tl/norm.py:75` `norm_tl`、`v4/norm.py` | 可复用。需要对照 `rstd_1ulp` 核对一致性。 |
| `swiglu.swiglu` :40 | 带 clamp 的 SwiGLU（up 限制在 ±10，gate 限制 ≤10）× 路由权重：gate/up `[R,I]`，weights 为 `(R,1)` 或 `(1,1)`，limit `(1,1)` → `[R,I]`。 | `tile_kernels.quant.swiglu_forward_and_per_token_cast`（`v4/mlp.py:19`） | 可复用。它不做 cast；fp8 舍入交给下一个 GEMM 的 QDQ。 |
| `engram_gate.gated_add` :41 | 对 4 条流中的每一条：rms(h)·rms(key) 归一化点积 → sign·sqrt → sigmoid 门控；输出 h + gate·value。h/key `[R,4*5120]`，value `[R,5120]`，weight `(1,4*5120)` = q_weight·k_weight。 | `tile_kernels.engram.engram_gate_fwd`（`v4/engram.py:155`） | 对照我们 fused-weight 的语义核对后可复用。 |
| `pp_permute.kernel` :395 | 用 `nki.collectives.collective_permute` 加 `nisa.core_barrier` 做 stage 之间的交接。**没有被调用**：CC 运行时拒绝了它的 replica group。 | — | 跳过。 |

**没有 MoE kernel 时，MoE GEMM 和反量化是怎么做的**（`model.py:2287-2380`、`:2851-2928`）：
1. router：`bf16_linear` + sqrt-softplus + nkilib top-k。
2. host 把 indices 拷回来，为每个本地专家建一个行列表，补齐到不小于 128 的 bucket。
3. 每个专家一张编译好的图：`index_select` 取行 → `fp4_linear_2core`（w1、w3）→ `swiglu` → `fp4_linear_2core`（w2）→ `index_add` 加回 fp32 累加器。
4. 全局 all-reduce。
5. 共享专家：同样的链路，换成 `fp8_linear_2core`，对所有行计算。

反量化在每次 GEMM 内部完成（权重字节 → private HBM 中的 bf16），不保留常驻的 bf16 副本。

## 4. 时间花在哪里，TPOT ≈ 120 ms 靠的是什么

PR 没有给出 profile 拆解。注释里能拼出这些数字：
- 图执行：每次固定约 1.12 ms，每次 forward 约 120 次。128 行专家 GEMM 的 roofline 约 154 µs，所以路由专家这部分时间大头是固定开销，而不是计算（`model.py:2353-2363`）。
- 全局 all-reduce：约 43 × 1.9 ms ≈ 每步 80 ms。
- decode 一步 ≈ 320-370 ms，处理 6 行。TPOT = 一步时间 / (1 + 接受数)，在 sonnet 上接受长度为 2.75（各位置平均接受率 36%）。
- prefill：每次 forward 的固定开销 F ≈ 0.3-0.8 s，外加每 token 约 1.06 ms。

起决定作用的设计：
1. **全部常驻 HBM**（每 rank 13.9 GiB）：forward 期间没有任何 host 流量。
2. **attention TP=8 并复制**，attention 的归约限制在 2 个 device 内。
3. **专家按 id 划分、token 复制**：不需要 dispatch，只要一次 all-reduce。
4. **GEMM 用满每个 LNC=2 核里的两个物理核**：先 fp4、再 fp8 的 2core 化把 TPOT 从 142.8 ms 降到 116.9 ms。
5. **每类算子都用 NKI**，减少图和算子的数量。
6. **prefill 分段大小 4096。**
7. **DSpark**（约 2.75 倍）。

他们实测过、但没有收益的尝试：在 SBUF 里融合反量化、共享压缩 KV 的读取、`NEURON_RT_IO_RING_CACHE_SIZE`、collective 拆成 8×8、`_EXPERT_ROW_FLOOR` 设到 128 以下。

## 5. 我们的集成计划（按预期收益排序）

1. **#9 静态图 + #12 放置（结构性改动，从约 150 s 降到 1 秒以内）。**
   - 每个 bucket 编译一张 decode 图，NKI 算子都放在图里。
   - 去掉所有 host 同步。为了避免 PR #71 那种每层一次的 `.cpu()`，**让 6 个本地专家都跑一遍 bucket 内的所有行**，没路由到该专家的行把路由权重设为 0。他们的 GEMM 本来就把 M 补到 128，所以 decode 时每个专家的指令量和他们 128 的 bucket 一样，而图保持静态。
   - bs1 下考虑**把 attention DP 换成 attention TP=8 组**：现在是一个 rank 算全部 64 个 head，其余 63 个 rank 在等。
   - token 复制以后，我们的 all-gather dispatch 就是空操作，combine 也变成一次全局 all-reduce，和他们的做法一样。需要实测 all-reduce 和 reduce-scatter 哪个更好。
   - 照搬他们 `_COMPILE_OPTIONS` 的编译参数。每次全局 collective 按约 1.9 ms 估算。
2. **#10 GEMM：复用 `blockquant`。**
   - `fp8_linear(_raw)_2core` 接到 `fp8_fp4_gemm_nt`/`fp8_einsum` 后面，`fp4_linear_2core` 接到 `m_grouped_fp8_fp4_gemm_nt_contiguous` 后面（在本地专家上循环），`bf16_linear`/`fp32w_linear` 接到 bf16 和 hc 的 GEMM 后面。
   - 权重保持 checkpoint 的 uint8 字节布局（这也符合我们"fp8 按字节搬运"的规则）。
   - 需要决定：接受 224 的 QDQ，还是把 kernel 改回 448 并重新验证饱和相关的编译参数。
   - 这是 kernel 层面收益最大的一项。
3. **#15 DSpark。** 他们的 TPOT 里含有 2.75 倍的放大。没有 DSpark 的话，我们一步要做到 117 ms 以内，而不是约 320 ms。要照着做的设计是在 forward 内部起草（从第 37-39 层抽取、5 token 的草稿块、Markov head）。
4. **#11 attention / indexer。**
   - 选定 KV 布局：bf16 cache 配 `gathered_attn_idx`，或者保留 fp8 记录、先 gather+反量化再调 `gathered_attn`。
   - `index_score` 接到 `fp8_fp4_(paged_)mqa_logits` 后面，`compress_pool` 接到 `tl/compress.py` 后面。
   - `deep_select.topk` 和 #22 借用 `vllm_neuron.functional.topk`（nkilib）。
5. **#14 Engram。** 照搬他们的放置：按行切分的 bf16（或 fp8，2.9 GiB）表放在 HBM，hash 和 gather 都在 device 上做，再做一次全局 all-reduce。这样 host gather 和 H2D 就完全去掉了。前提是 HBM 有余量（#21）：他们的稠密部分每 rank 只有 3.9 GiB，是因为 attention 用了 TP=8、词表切了 64 份。再加上 `engram_gate`。
6. **#13 小算子。** `hc_mixes`、`hc_expand`、`hc_collapse`、`hc_attn_out`、`rmsnorm`、`swiglu` 单个都不大，但合起来能消掉每层大部分算子数。放在 #9 之后做，这样 profile 才是真实的。

**在我们的栈上复用任何 kernel 之前要先确认：**
- `torch_neuronx.nki_op`（2.14.3）是否接受*返回* `shared_hbm` 输出、使用 `private_hbm` 临时空间的 `@nki.jit` kernel，能否同时在 eager 和 `torch.compile(backend="neuron")` 下运行，能否在 LNC=2 下以 grid=2 发射。
- 间接 DMA（`.ap(vector_offset=…)`）经过 `nki_op` 能否工作。
- 形状断言（k%128、K%128、heads≤128）。
- 与我们 CPU shim 的数值一致性，误差要在量化噪声以内。

**许可证。**
- vllm-neuron 是 Apache-2.0（`LICENSE`；`NOTICE`："Copyright Amazon.com, Inc. or its affiliates"）。DeepInfer 也是 Apache-2.0，两者兼容。
- kernel 文件本身没有逐文件的许可证头。只有 `kernels/__init__.py`、`model.py`、`weight_loader.py` 等文件带 `SPDX-License-Identifier: Apache-2.0`。
- 拷贝时：给每个拷过来的文件加 SPDX 头和来源说明（仓库、PR #71、SHA `30d54d1`、作者 Zifan He / OswaldHe）；改过的文件要注明已修改（Apache-2.0 §4b）；把相关的 NOTICE 文本带进我们的 NOTICE 或第三方声明。
- PR 还没合并，所以要钉住 SHA，合并后再核对一次。

</details>