An fp8 tuple input fell back to the deep_gemm shim, which dequantized the whole fp8 weight in the graph on
every call (Engram fused_kv_projection: 9 ms of an Engram layer's prefill graph, GpSimd POW for the scales).
Now decoded to bf16 and run through the NKI kernel. TTFT 1.71 -> 1.66 s (bucket 512), tokens unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
18 of 64 ranks compute a 128-column slice of the 2304-wide shared expert (the rest add zeros, same graph);
the partial joins the routed experts' sum before the existing all-reduce. Decode step 591 -> 581 ms.
Summation order changes, so greedy outputs differ from before (coherent; tokens/step 2.70/3.10/3.34,
TPOT p50 184 ms). Tiny golden 9/9. Also DEEPINFER_NEURON_HOST_TIMING (host time per region call).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The fp4 expert GEMMs dequantize the whole weight to bf16 in HBM on every call
(device profile: the MoE graph reads 756 MB a layer), and a rank computes all
6 of its experts while a 6-row step routes to about one. fp4_linear_2core_gated
takes a per-expert flag (any row routed to it) and runs the dequant and GEMM in
an nl.fori_loop bounded by it, so an unrouted expert costs a zero fill. Decode
sizes only (<= 16 rows); prompt chunks route to every expert anyway.
bench70 (vllm-sonnet 512, TP=64, nextn=5): TPOT 258.9 -> 248.4 ms, step
738 -> 711 ms, tokens per step identical. Tiny cp=2 nextn=5 on device: 9/9
golden tokens, drafts unchanged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
At LNC=2 a logical NeuronCore is two physical cores and a kernel launched at
grid 1 only has code for one of them. The vllm-neuron#71 _2core kernels split
the output columns by program_id; launched at grid 2 they are 1.7-1.9x faster
on every decode GEMM shape and bitwise equal on one core (fp8 1x1280x36864
2.22 -> 1.17 ms, fp4 64x5120x4608 0.68 -> 0.40 ms). Used when n % 256 == 0;
DEEPINFER_NEURON_BQ_2CORE=0 turns it off.
Device profile of a decode attention region (neuron-explorer): 12.9 ms, of
which the q_b_proj fp8 GEMM alone is ~1.8 ms of vector-engine dequant.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Layer-major KV / state layout so per-layer views are contiguous graph inputs
(the executor copied non-contiguous inputs every call).
- One compiled entry per (class, method), rank as a registered buffer, no
layer_idx guard: every layer reuses the same graphs and loaded NEFFs.
- fp8 weight rescale on the host (numpy LUT, ~2 min instead of 35+ min of
eager per-shape compiles); engram sf packing on the host.
- int16 ue8m0 packing in fp32 (int32 interleaved slices narrowed to int16
crash neuronx-cc, NCC_IBCG902).
- Token embedding and the text model's vision tower stay in host memory on
Neuron (2.5 GiB of HBM a rank).
- offline example: --max_num_seqs.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Decode runs eagerly and compiles per-layer regions (attention half, MoE) with stage boundaries
(neuron_ops/regions.py, DEEPINFER_NEURON_GRAPH_MODE=regions); layer KV views are sliced from the
roots inside the graph (_LazyKVCache).
- vllm-neuron#71 NKI kernels (Apache-2.0, vendored under neuron_ops/nki/ with attribution), called via
torch_neuronx.nki_hop.wrap_nki: blockquant (fp8/fp4 GEMMs incl. grouped einsum, MoE and shared
expert on bf16 activations), sparse_attn / index_score / compress_pool (nki_attention.py),
hc_mixes / hc_expand / hc_collapse / hc_attn_out / rmsnorm / swiglu (nki_hc.py, nki_norm.py).
- Compile-safe formulations: row-major scales, no dtype bitcasts of graph inputs (uint8 weight
tensors), whole-row KV stores, byte assembly in int32 narrowed once at the store, arithmetic int16
scale packing, plain fma/rsqrt and bf16 F.linear under compile; default backend deepinfer_neuron
(torch fp8 carried as uint8, avoiding NCC_EOCP001 with the kernels' OCP fp8).
- Host: tiny model 9/9 vs golden on all prompts; device: kernel tests pass (13 attention, 88 hc/norm);
region decode graph compiles except the attention projection stage (being reduced).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
