DeepInfer/docs/neuron/vllm_neuron_pr71_analysis.md
whn09 02e099fff8 Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24)
- 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>
2026-10-22 11:59:37 +08:00

45 KiB
Raw Blame History

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

Source: vllm-project/vllm-neuron#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.
中文

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

来源:vllm-project/vllm-neuron#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,合并后再核对一次。