Within 1 bf16 ulp of the gathered kernel on the full model (mask equality checked on every layer). With a 1024
context bucket: TTFT 1190 -> 1113 ms; tokens differ from the kernel's by rounding (2.89 -> 3.02 tokens a step).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_dense_attention matches gathered_attn_idx to ~1 bf16 ulp (same error vs an fp64 reference). It only applies to
tables <= DEEPINFER_NEURON_DENSE_ATTN_MAX_KV rows, which a prefill chunk's whole-pool table exceeds; bounding that
table (DEEPINFER_NEURON_PREFILL_TABLE_BOUNDS) went through the compacting gather and doubled TTFT, so it is off.
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 table path decoded the whole KV pool every layer (160 blocks x 128 slots);
the fp8 record decode is integer work on GpSimd and took 6-10 ms a layer
(device profile of the 6-row attention graph, bench47). A decode step's
requests reference at most ceil(ctx / page) + 1 blocks of a scope each: those
blocks are chosen with cumsum + scatter (static shapes), gathered whole (one
contiguous run each; a per-slot gather was slower than the full decode) and the
indices remapped into them. Sizes are rounded to a power of two (neuronx-cc ISA
validation 'same_element_count_t2d' otherwise). Prompt-chunk buckets keep the
whole pool. The graph runner sets the bounds (requests, context bucket) per
replay.
bench50 (512 in, 128 out, TP=64, nextn=5): TPOT 234.1 -> 208.4 ms, step
907 -> 805 ms, per-request tokens-per-step identical to bench41 (3.85, 4.38,
1.98). Tiny cp=2 nextn=5 with a 128-block pool: 9/9 golden tokens.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A 128-row chunk's attention ran 122 ms a layer as one graph and 62 ms split at
the stage boundaries (bench3 vs bench4); decode's 1-row graphs are faster
whole. split() now cuts when DEEPINFER_NEURON_REGION_SPLIT (default off now)
or, under DEEPINFER_NEURON_REGION_SPLIT_PREFILL (default on), when the step
has more than one row. Syncing after every region call (the queueing
hypothesis) did not help: 37.7 s TTFT at 512.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The op census of a decode step (DEEPINFER_NEURON_REGION_TIMING) found 85 eager
aten ops in the forward outside the region graphs but 1905 in sampling (the
top-p/top-k bisection and the int64 philox draw), ~0.6 s of NEFF launches.
- The log-softmax / top-p / top-k part runs as one compiled graph: 1.4 ms
instead of 46.7 ms on one core for the [1, 4096] force_topk row, logprobs
within 1e-9.
- Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so
the two-level draw returns their argmax; an all-greedy batch takes the argmax
from the graph and skips the draw. The philox draw stays eager for sampled
rows: neuronx-cc graphs reject 64-bit constants outside the int32 range
(NCC_ESFH001), an int32 philox is #3.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- The state buffer's column regions get their own allocations on Neuron:
a layer's uint8 SWA cache and fp32 compressor state are both region graph
inputs, and aot_autograd compares offsets of differently typed views of one
storage in their own element units (false overlap -> 'input mutations on
views with different dtypes').
- paged_mqa_logits chunk path: gather the keys the chunk's last token sees
(lens.max()), not the first token's.
- Debug: DEEPINFER_NEURON_ALIAS_DEBUG prints aliased graph inputs; a cProfile
of decode step 8 under DEEPINFER_NEURON_REGION_TIMING.
Tiny V4.1 on a NeuronCore, bucket 128: short 8/9 (= eager), long 715 tokens
in 6 chunks 9/9 (eager 8/9) in 26 s instead of 978 s of eager prefill.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- PrepareArgs.prompt_as_decode: decode-layout metadata for a prompt chunk;
input_ids are the prompt's own tokens instead of last token + MTP drafts.
- DEEPINFER_NEURON_PREFILL_BUCKET adds one graph size; the model runner pads
a prompt chunk to it with the decode placeholders (every EP rank takes the
same bucket, so the static collectives match). Sampling still sees a prompt
step.
- nki_attention.paged_mqa_logits: one gather + one kernel call for a chunk
(one sequence) instead of a T-way loop.
- bf16 linear for the LM head eagerly on a NeuronCore (b.float() was a
2.6 GiB copy a call); DEEPINFER_NEURON_REGION_TIMING; --num_prompts.
Host mode, tiny V4.1, chunk 128: same tokens as eager prefill (short 9/9,
long 715 tokens in 6 chunks 8/9 against golden).
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>
- 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>
