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>
Both drain the device stream in torch-neuronx (a cast round-trips through the
host, a scalar operand is a host->device copy): int32 ids index directly and
out-of-range rows are cleared with masked_fill. Tiny cp=2 nextn=5 on device:
9/9 golden tokens.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- DFlash._greedy_drafts: post norm, LM head, Markov bias, confidence head and
the argmax of the 5 draft steps as one region graph when every row is greedy
with no top-p/top-k filter.
- DEEPINFER_NEURON_HOST_EMBED defaults off: in torch-neuronx every host<->device
copy drains the stream, and the host lookup (D2H ids, H2D rows) left the cores
idle for 50-100 ms gaps between the draft graphs (runtime trace bench36).
- Eager device embedding lookups are one F.embedding with a validity mask (no
host scalar operand).
- bench34 / bench37 (512 in, 128 out, TP=64, nextn=5): TPOT 268.2 -> 264.0 ->
234.7 ms, step 1025 -> 907 ms, TTFT 5.44 -> 5.01 s.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- fused_sample._host_flags: the top-p/top-k and greedy flags of the step's
sampling parameters are read from the device once a step (a DSpark step
samples the same tensors six times); reset_host_flags at each sample step.
- NEURON_RT_INSPECT_* stays on the driver (64 ranks would overwrite one
output directory).
- VocabParallelEmbedding.ids_in_range: one F.embedding instead of the
zero-filling gather (~5 eager NEFF launches).
- bench33 (512 in, TP=64, nextn=5): TPOT 272.6 -> 268.2 ms.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
neuronx-cc segfaults (LegalizeSundaAccess) when the 6-row projections are merged
with the indexer, so that boundary stays; the other fine splits are prompt-only.
bench28 (512 in, 128 out, TP=64, nextn=5): TPOT 277.4 -> 272.6 ms, 3.85 tokens
a step.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The table path at 6 rows computed wrong attention (bench21: acceptance 1.27);
padded to 8 rows it is right and the step drops from ~1.8 s to ~1.07 s.
bench23 (512 in, 128 out, TP=64, nextn=5): TPOT 466.4 -> 277.4 ms, 3.85 tokens
a step, per-position acceptance identical to bench18.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
bench18 (512 in, 128 out, TP=64, nextn=5): acceptance 3.85 tokens a step,
TPOT 465.9 ms. bench21 with the table path for 6-row steps: acceptance 1.27,
TPOT 848 ms, so verification keeps the per-row path. Whole-layer 6-row graphs
crash neuronx-cc on the full model (DEEPINFER_NEURON_SPLIT_MIN_ROWS, default 1
= split as before).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- DFlash forward_drafts per layer and update_main_cache (projection, then each
layer's cache write) run as region graphs: eagerly the fp8 cast of the
projection norm does not compile.
- neuronx-cc rejects the sparse decode attention at 6 rows (a DSpark step with
5 drafts; ISA validation 'same_element_count_t2d'), 4 and 8 rows compile:
multi-row decode steps pad to a power of two with empty rows.
- bench_neuron.py --nextn: DSpark drafts, reports tokens per decode step.
- Tiny checkpoint on device, cp=2, free-running greedy with drafts: nextn=3,
5, 7 all 9/9 golden tokens (short), nextn=5 long 9/9.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- fused_sample reject mode for greedy rows (accept iff the draft is the
argmax), get_accept_result / finalize_verification (static final rows),
store_mtp_info (static pool writes, slot 0 kept), confidence_head_fused,
dflash_input_store, get_fused_dflash_get_input: the frontend refs with the
kernels' argument lists.
- CP followers run the draft forward (its MoE has collectives); the driver's
accepted offsets reach them through one all-reduce on Neuron (no eager
broadcast), the request first rows through the graph metadata.
- run_deepinfer_tiny.py --nextn: free-running greedy with drafts. Host tiny,
nextn=5, eager: short 9/9 golden tokens; long matches up to its known near
tie at step 5.
- bench17: the two-level MoE all-reduce does not move TPOT (433 ms).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Microbenchmark at 64 ranks, [1, 5120] fp32 in a compiled graph: world
all-reduce 3.37 ms, 8 consecutive ranks 0.41 ms, 4-core device group then a
16-device group (one core each) 2.01 ms; a strided 8-device group is rejected
by the runtime. The replicated-token MoE all-reduce (one a layer) now takes
the two-level form (DEEPINFER_NEURON_HIER_ALLREDUCE=0 restores one world
all-reduce). Host tiny cp=8 matches golden.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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>
Attention is head-sharded within lanes of 8 consecutive ranks of the CP group
and replicated across lanes, as vllm-neuron#71 does: a 64-rank all-reduce is
~1.9 ms of latency against ~0.1 ms for 8 ranks on two adjacent devices. With
64 heads and 8 groups each rank holds exactly one wv group. Host tiny cp=8
with lanes of 2: golden tokens up to near ties.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
For a prefill chunk the per-row path decoded rows x topk KV records (128 x 640
a layer: 31.8 ms of a 49.6 ms attention region at TP=64). The chunk path now
decodes each cache scope once into a bf16 table (pages x page slots, NaN-free)
and the gathered_attn_idx kernel gathers each row's slots in SBUF, as
vllm-neuron#71 does. Used while the caches hold <= 32768 slots
(DEEPINFER_NEURON_ATTN_TABLE_SLOTS); size the pool with --num_kv_blocks.
Tiny V4.1 on a NeuronCore: short 8/9, long 9/9, cos >= 0.9979.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The name came from the parent pid, which CP followers do not share with their
driver (they are its children): followers created a second table, never
unlinked (622 GB of /dev/shm leaked over five runs, then a host OOM), and each
side read zeros for the other's shards. Rank 0 now publishes the path through
the distributed store, and opening asserts the size so it can never create a
file.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Its periodic load-statistics all-gather spans the EP group; the driver waited
in it forever (allgather seq=10 timeout after the warmup request at cp=64).
Also bench_neuron.py --cp.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With a CP group running one replicated batch (DEEPINFER_NEURON_TP_ATTN, on
when cp > 1), each rank keeps its heads' rows of q_b, its groups' slice of wv
(or its heads' columns of one group), the matching out_proj columns and its
attention-sink entries; _wv_and_o_proj all-reduces the partial outputs over
the group (exact: wv and out_proj are linear in the heads; asserts no clamp).
Q-a / KV / indexer / compressor stay replicated.
Host tiny (8 heads, 2 groups): cp=2 (1 group a rank), cp=4 (2 heads of a
group), cp=8 (1 head) all give the golden tokens up to near ties.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Followers claim core cp_rank (DEEPINFER_NEURON_LOCAL_RANK at spawn) and get
the graph's host metadata through shm (no in-graph broadcast on Neuron),
including the prompt-as-decode bucket and the context limit; they skip the
sampler. max_num_prefilling_tokens is not divided by cp on Neuron (prompt
chunks run replicated).
- DEEPINFER_NEURON_NUM_KV_BLOCKS: the fixed KV pool as an env var, so spawned
CP followers take it too.
- static MoE: when the EP group is the CP group the tokens are already
replicated: no all-gather, all-reduce instead of reduce-scatter (64x fewer
expert rows a rank at cp=ep=64).
- run_deepinfer_tiny.py --cp. Host tiny: cp=2 and cp=4 match the cp=1 tokens
(short 8/9 at a near tie, long 9/9).
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>
DEEPINFER_NEURON_CTX_BUCKETS: the model runner sets the step's context limit
to the smallest bucket that holds every sequence, and the dense kv_ca_index /
indexer logits are sized to it instead of max_model_len (one more set of
attention graphs a bucket). Prefill cost scaled with the limit: a 128-token
chunk ~9 s at 16384 against ~5.3 s at 4352. Host tiny check: same tokens as
without buckets up to near ties.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Only dense compressed-KV scopes need the shared-gather torch path; layers whose
scopes are all per row (top-k + SWA, SWA only) take the sparse_attn kernel as
in decode. bench4: the torch path was 41.6 ms a layer for a 128-token chunk.
Tiny V4.1 on a NeuronCore: short 8/9, long 9/9, logit cos >= 0.9977.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A 128-token chunk at EP=64: forward_capacity 108.5 ms a layer against 58.1 ms
for the dense MoE, and with capacity 1024 every layer overflowed and reran
densely. The chunk's attention (122 ms a layer) is the larger prefill cost.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The overflow decision was rank-local, and the dense rerun has collectives of
its own: one rank rerunning left the others waiting in its all-gather (NRT
'Collective operation pending'). The flag is now max-all-reduced over the EP
group. The bench gave every rank the same prompt, which routes 64 copies of
each token to the same experts (overflow on every layer); real requests
differ, so the prompts now take the rank in their seed.
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>
With runtime_max_model_len <= candidate_topk_blocks * candidate_block_size
every consumer already skips the candidate filter, so the source layer's block
top-k feeds nothing; below 16384 positions it is not even defined (fewer
blocks than candidate_topk_blocks, assert in _topk_candidate_block).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The bisection is 32 rounds over the vocabulary: ~0.5 s a decode step at 129280
even compiled, and the default parameters (top_p 1, top_k -1) filter nothing.
Such batches take the plain temperature log-softmax: 529.6 -> 1.3 ms on one
core, same tokens, logprobs within 1e-9.
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>
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>
torch.remainder ran ~140 ms a call eagerly on the device and was 554 of the
640 ms a decode step spent sampling. a - (a // b) * b gives the same result for
the non-negative operands (checked on 100k random values).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- deep_select.topk on a NeuronCore with k >= 1024 (the sampler's force_topk
4096 over the 129280 vocab) uses nkilib's rotational_topk: torch.topk with
that k does not compile on the device (NCC_IXCG857). 8.4 ms on one core,
values identical to torch.topk.
- EngramBuffer on Neuron: the ranks of a trn2 instance share one host, so the
table shards are written into /dev/shm-backed tables (unlinked once mapped)
and every rank gathers its own rows: no all-to-all and no host sync on split
sizes per fetch. DEEPINFER_NEURON_ENGRAM_SHARED=0 restores the sharded path.
Host EP=2 tiny: logits bitwise equal to the sharded path.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- When runtime_max_model_len <= candidate_topk_blocks * candidate_block_size,
the candidate blocks hold every position, so the filter keeps them all and
the dense indexer logits give the same top-k. On Neuron the candidate path
gathers all candidate keys per query token: [128, 16384] rows a layer for a
prompt chunk, the NEFF that still ran HBM out in rg16.
- static_moe_forward takes the decode metadata's valid_token_mask: padding rows
select no expert.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- nki_attention._chunk_attention: a prompt chunk run as a decode step gathers a
dense compressed-KV scope once (every token's list is a prefix of the
longest) and masks it per row; top-k and SWA scopes stay per row. The
per-row gather was [128, 8192, 512] for a ratio-2 layer of the full model:
3.8 GB moved per call and a 3 GB NEFF (model code) that ran HBM out.
- decode_part_timer: the decode timers are no-ops inside region graphs. Their
f-string labels read layer_idx, so Dynamo guarded on it and retraced every
layer instead of every layer kind.
Tiny V4.1 on a NeuronCore, bucket 128: short 9/9, long 9/9 (logit cos >= 0.997).
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>
Fixes found by reducing neuronx-cc NCC_IINL902 crashes to minimal reproducers (crash rates over repeated
compiles, the compiler is nondeterministic on these graphs):
- uint8 byte assembly of narrowed values (KV records, e2m1/e4m3 codes, ue8m0 exponents) crashed or
silently corrupted bytes: producers keep int32 codes under compile and stores narrow once
(fused_transform, kv_quant, store_kv_cache; neuron_compile.encode_codes).
- int16 packed scales written through a uint8 view and read back in the same graph: the Neuron kernel
now receives the int16 storage and packs byte pairs arithmetically (_alloc_quant, return_quant).
- fp8 legalization pass treated plain float->bf16 conversions as fp8 carriers (wrong values + crash).
- Projection stage split further (nullcontext instead of a disabled CP session).
Device (tiny model, regions mode): every decode graph compiles; argmax vs golden short 9/9,
medium 8/9 (step 6 near-tie, margin 0.065), long 9/9, cos 0.997-0.9995.
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>
Uneven all_to_all on the neuron process group is unreliable at scale (EP=64:
"Invalid NEFF" in prefill; test_ep at 64 ranks hung after 9/18 modes).
EPBuffer now moves tokens without all_to_all:
- dispatch: one all_gather of x / sf / weights / topk_idx packed as uint8
rows (row count = global max via a 0-dim MAX all_reduce, or the static
num_max_tokens_per_rank with DEEPINFER_NEURON_EP_AG_ROWS=cap); each rank
selects the rows routed to its local experts in the same (source rank,
source token) order as before, so layout / metadata / cached dispatch are
unchanged.
- combine: per-token fp32 partials over local experts (bf16-rounded as
before) + weights, one fp32 SUM reduce_scatter (all_reduce fallback with
DEEPINFER_NEURON_EP_REDUCE=ar), chunked under
DEEPINFER_NEURON_EP_CHUNK_BYTES (64 MiB).
- The a2a path stays behind DEEPINFER_NEURON_EP_ALGO=a2a.
test_ep passes on the neuron group at 2/4/8/16/32/64 ranks (combine
bit-identical to the reference in practice); new test_collectives.py found
that int32 all_gather with full-range values is wrong on a 16-rank neuron
group for >= 1 MiB outputs (uint8 / fp32 fine) -- payloads move as bytes.
Tiny model EP=8 on NeuronCores: 25/27 argmax vs the reference, same as EP=1.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The DP-wide prefill vote runs every DEEPINFER_PREFILL_ATTEMPT_INTERVAL (70) /
..._ANY_ (100) decode steps. Requests that become schedulable after step 0
(multimodal-processor round trip) then wait 70 placeholder steps; with eager
execution a full-model step on the NeuronCore takes ~150 s, so the smoke test
sat ~3 hours in placeholder steps. Default both intervals to 1 on Neuron.
offline_inference_neuron: force-topk off (#22), 400-block KV pool.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- requirements-neuron.txt: requirements.txt + dev deps without tilelang /
tile-kernels (pure-torch drop-ins are vendored), torch 2.14 +
torch-neuronx 2.14 from the Neuron repository.
- build_vendor.sh: nothing to build for neuron; setup.py drops tilelang /
tile-kernels from install_requires on neuron.
- setup.py and get_platform(): an exported DEEPINFER_ACCELERATOR wins over
device detection, so images build where no accelerator is visible.
- docker/neuron/Dockerfile (+ per-Dockerfile ignore): on the native PyTorch
Neuron DLC public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-
sdk2.33.0-ubuntu26.04; build.sh with TORCH_DEVICE_BACKEND_AUTOLOAD=0 and an
import self-check in Neuron CPU mode.
- README: Trainium in supported hardware and a "DeepInfer on AWS Trainium"
section (English + folded Chinese); the pip torch-neuronx wheel needs
glibc >= 2.43, so on Ubuntu 24.04 layer the venv on the DLAMI env.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Device-memory fixes found bringing up the released checkpoint (16.45 GiB of
weights per 24 GB core):
- Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the
all-gather of the replicated SF table (~3 GB per layer, allocated on the
device -> "Failed to allocate resource").
- gather_tl: index_select / advanced indexing of an int8 table allocates
several times the table on the device (+4.9 GiB for the 1.3 GB embedding);
gather through F.embedding on an int32 / fp16 view instead (bit-exact, no
extra memory).
- deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements
and one expert at a time instead of whole projections / all local experts
in fp32 (OOM on a fragmented pool); bit-identical results.
- transfer_to_neuron: import torch.utils._device explicitly (missing when
torch_neuronx is not auto-loaded).
examples/neuron/offline_inference_neuron.py: offline inference in basic mode
(text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for
the candidate-block filter), one process per NeuronCore.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Three torch-neuronx 2.14 device behaviours, each worked around in a
bit-identical way:
- torch.cat of float8 tensors is wrong on the device: merge_linear_weights
(fused q_a/kv, q_b, ... projections) now concatenates the raw bytes, like
the HF loader's read_concat.
- int64 constants are narrowed to int32 by the compiler
(--implicit-integer-downcast): deep_select.topk used int64.max as the sort
sentinel, which became -1 and pushed the selected indices out of rows with
end < topk -- the first compressed layer then attended to no compressed KV
for short prompts. The sentinel is now vocab_size.
- advanced indexing of non-contiguous float8 views is wrong: the remaining
site (deep_gemm sparse MQA logits) goes through uint8 views; deep_ep's
byte path also covers e8m0.
DEEPINFER_NEURON_FP8_GUARD=1|raise (neuron_ops/fp8_guard.py) reports every
cat/stack/advanced-index/gather/scatter/where on a float8 tensor with its call
site; the tiny end-to-end run and all CPU unit tests report zero sites.
Tiny checkpoint on one NeuronCore, teacher-forced vs the official reference:
short 9/9 (min cos 0.9985), medium 8/9 (0.9991), long 8/9 (0.9978); logits
rel <= 0.066, within the reference's own noise floor -- same as CPU mode.
Also: probe_layer0 PROBE_LAYERS=N and attention-input capture, device_diff
copies tensor attributes (replays deep_gemm sparse MQA metadata), device
primitive checks for the narrowing and strided fp8 casts.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- deep_ep comm: torch-neuronx picks the all_to_all lowering per rank from
that rank's own splits (even / NKI a2a-v / 1-D-only padding fallback), so
ranks could diverge. All-gather the split matrix first and choose one path
for all ranks: even a2a, NKI a2a-v for groups of 4/8, otherwise a padded
even a2a (2 ranks, 16 ranks where a2a-v fails on runtime 2.35).
- PlatformNeuron: LOCAL_RANK=0 / LOCAL_WORLD_SIZE=1 so every one-core
process is neuron:0 of a 1-device local world (otherwise rank > 0 died in
_lazy_init at EP/DP > 1).
- tests/neuron/deep_ep: references on host copies with the gloo group,
buffers on the device; ep/engram/bucket pass at world size 2/4/8/16 on the
neuron process group, and 1/2/4 on gloo.
- run_deepinfer_tiny.py --ep N: one process per rank / NeuronCore; tiny
model EP=2 runs end to end (CPU: within the noise floor).
- run_all.sh: per-core busy check, bucket + world size 16 device stages.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
torch.cat of float8 tensors (e4m3 and e8m0) returns wrong data on the
neuron device (torch-neuronx 2.14). The HF loader concatenates w1/w3 (and
their UE8M0 scales) into the fused gate_up weights, so on the NeuronCore
every routed and shared expert gate_up scale was zero and the MoE output
vanished (issue #1). Concatenate 1-byte float dtypes through uint8 views,
which is bit-identical on every platform.
tests/neuron/device: create the neuron process group before the runtime
initializes (torch-neuronx requirement) and add the float8 cat checks.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Engram: gather rows where the table lives (host by default, read via UVA
on GPU) and move them to the compute device.
- EPHandle: the topk_idx version guard tolerates inference-mode tensors on
the neuron device that do not report is_inference().
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pure-torch sparse attention followed FlashMLA's tests/ref.py (fp32 P).
The CUDA kernels (sm100 prefill phase1.cuh:238-253) and the official
reference round P = exp(s - running_max) to bf16 before the P@V GEMM in an
online softmax; do the same with 64-key blocks.
Localized with examples/neuron/probe_layer0.py (per-op reference vs DeepInfer
capture): layer 0 is now bit-exact up to two rounding points that are
DeepInfer's own GPU/Ascend behaviour (SwiGLU -> FP8 from fp32, bf16 routed
partial in the MoE combine). CPU end-to-end on the tiny checkpoint: 26/27
greedy tokens match the reference, the miss being a 0.0096-margin step the
reference flips against itself.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout,
pure-torch replacements for the official reference kernels, golden
generation, a teacher-forced DeepInfer runner (LLM offline API) and a
comparison tool. CPU mode: 25/27 greedy tokens match, worst logits
cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060.
- model_loader: check parameters against the platform device, not is_cuda.
- deep_gemm paged MQA logits: do not dereference block-table slots past the
context.
- bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would
serialize host compute).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
