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>
- 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>
- 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>
