Attention half and MoE as one region graph per decoder layer: half the NEFF
launches a step (the attention lane all-reduce did not move the decode step,
447.8 ms, so per-launch overhead is the next suspect). Host tiny cp=8 matches.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- serve/: the minimal OpenAI-compatible server of deepseek-ai/DeepInfer#147,
with a Neuron worker setup (basic mode env, per-rank NeuronCore, fixed KV
pool via --num-kv-blocks, no MTP lookahead) and a /v1/completions endpoint
(prompt text or token ids, no chat template): what vllm bench serve
--backend openai drives, as vllm-neuron#71's run_bench.py does.
- examples/neuron/bench_neuron.py: the same TTFT/TPOT matrix offline.
- static MoE: one packed fp32 all-gather instead of three (each 64-rank
collective is ~2 ms of fixed latency a layer); a capacity-grouped expert
path for prefill chunks (DEEPINFER_NEURON_MOE_CAPACITY, dense rerun on
overflow; not yet validated on device).
- Sub-region timings at split() boundaries under DEEPINFER_NEURON_REGION_TIMING.
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>
