- 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>
45 KiB
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_xlaor NxDI. It useslibtorch_neuronx_lite(wrap_nkifor NKI,model.py:401),torch.compile(backend=vllm_neuron.envs.get_compile_backend_name())(model.py:445,:2234),torch.ops._c10d_functionalcollectives (model.py:225-240) andvllm_neuron.*helpers. Caveat: the numbers were measured on vllm-neuron 0.24.0.1.1.0, and the README liststorch/torch-xla2.11 in that stack. In 0.24 the default compile route isneuron_libtorch(lite on top of XLA); the native route needsVLLM_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
nki0.6 (import nki,nki.isakeyworddst=form,@nki.jit), the samenki>=0.6that ourrequirements-neuron.txtpins. There is no MoE kernel. Expert GEMMs areblockquant.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=1givesfullgraph=False,__init__.py:34). The model also creates its own compiled entry points withtorch.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.forwardis@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). - one per NKI kernel (
-
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_tokensis 4096 at 64 ranks (__init__.py:451-561). TTFT fitsceil(tokens/segment) × F + tokens × ~1.06 ms, with F ≈ 325-783 ms depending on segment and context. - Decode has
(1+5) × num_seqsrows.MAX_NUM_SEQS=1because 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.
- Prefill is segmented:
-
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 compressorpool_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).
- The window KV is vLLM's paged allocation, one per layer,
-
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 inself.last_draft, andDSparkProposer.propose(drafter.py) hands them to vLLM. There is no separate draft NEFF or warm-up. This needs the plugin's proposer registry (thevllm_neuron/vllm/spec_decodeand 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()onslot_mappingandblock_table(model.py:4766,:4831,:4943). int(start_pos)is read inBlock.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 throughDEEPSEEK_V41_EXPERT_FUSE, which is off by default.
- Each MoE layer copies
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_hbmand returned (functional style). Scratch space isnl.private_hbm. - 128-row tiles with a tail; reductions accumulate in fp32.
- Called through
wrap_nki(kernel)(libtorch_neuronx_lite.nki.nki_hop) insidetorch.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):
- Router:
bf16_linear+ sqrt-softplus + nkilib top-k. - The host copies the indices back and builds a per-local-expert row list padded to a bucket of at least 128.
- One compiled graph per expert:
index_selectrows →fp4_linear_2core(w1, w3) →swiglu→fp4_linear_2core(w2) →index_addinto the fp32 carry. - World all-reduce.
- 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:
- Everything resident in HBM (13.9 GiB per rank): no host traffic during the forward.
- Attention TP=8, replicated, so the attention reduction stays within 2 devices.
- Experts by id with replicated tokens: no dispatch, one all-reduce.
- 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.
- NKI for every op family, to cut the number of graphs and ops.
- Prefill segment 4096.
- 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)
- #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_OPTIONSflags. Budget about 1.9 ms per world collective.
- #10 GEMM: reuse
blockquant.- Put
fp8_linear(_raw)_2corebehindfp8_fp4_gemm_nt/fp8_einsum,fp4_linear_2corebehindm_grouped_fp8_fp4_gemm_nt_contiguous(loop over local experts), andbf16_linear/fp32w_linearbehind 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.
- Put
- #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).
- #11 attention/indexer.
- Choose the KV layout: bf16 cache plus
gathered_attn_idx, or keep fp8 records with gather+dequant plusgathered_attn. - Put
index_scorebehindfp8_fp4_(paged_)mqa_logits, andcompress_poolbehindtl/compress.py. - Borrow
vllm_neuron.functional.topk(nkilib) fordeep_select.topkand #22.
- Choose the KV layout: bf16 cache plus
- #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. - #13 small ops.
hc_mixes,hc_expand,hc_collapse,hc_attn_out,rmsnormandswigluare 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.jitkernels that returnshared_hbmoutputs, useprivate_hbmscratch, run under both eager andtorch.compile(backend="neuron"), and launch with grid=2 at LNC=2. - Whether indirect DMA (
.ap(vector_offset=…)) works throughnki_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.pyand similar files carrySPDX-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_functionalcollective(model.py:225-240),以及vllm_neuron.*工具函数。需要说明的是:性能数字测于 vllm-neuron 0.24.0.1.1.0,README 列出的这套环境里有torch/torch-xla2.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,全部基于
nki0.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)。 - 每个 NKI kernel 一个(
-
编译参数(
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 倍。
- prefill 分段:64 rank 时
-
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)。
- 窗口 KV 用的是 vLLM 的分页分配,每层一份,
-
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 个专家合进一张图,默认关闭。
- 每个 MoE 层都把
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):
- router:
bf16_linear+ sqrt-softplus + nkilib top-k。 - host 把 indices 拷回来,为每个本地专家建一个行列表,补齐到不小于 128 的 bucket。
- 每个专家一张编译好的图:
index_select取行 →fp4_linear_2core(w1、w3)→swiglu→fp4_linear_2core(w2)→index_add加回 fp32 累加器。 - 全局 all-reduce。
- 共享专家:同样的链路,换成
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。
起决定作用的设计:
- 全部常驻 HBM(每 rank 13.9 GiB):forward 期间没有任何 host 流量。
- attention TP=8 并复制,attention 的归约限制在 2 个 device 内。
- 专家按 id 划分、token 复制:不需要 dispatch,只要一次 all-reduce。
- GEMM 用满每个 LNC=2 核里的两个物理核:先 fp4、再 fp8 的 2core 化把 TPOT 从 142.8 ms 降到 116.9 ms。
- 每类算子都用 NKI,减少图和算子的数量。
- prefill 分段大小 4096。
- DSpark(约 2.75 倍)。
他们实测过、但没有收益的尝试:在 SBUF 里融合反量化、共享压缩 KV 的读取、NEURON_RT_IO_RING_CACHE_SIZE、collective 拆成 8×8、_EXPERT_ROW_FLOOR 设到 128 以下。
5. 我们的集成计划(按预期收益排序)
- #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 估算。
- #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 层面收益最大的一项。
- #15 DSpark。 他们的 TPOT 里含有 2.75 倍的放大。没有 DSpark 的话,我们一步要做到 117 ms 以内,而不是约 320 ms。要照着做的设计是在 forward 内部起草(从第 37-39 层抽取、5 token 的草稿块、Markov head)。
- #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)。
- 选定 KV 布局:bf16 cache 配
- #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。 - #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.jitkernel,能否同时在 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,合并后再核对一次。