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>
DEEPINFER_NEURON_MOE_FLAG_FORCE=0|1 forces the gated experts' flags (timing only, wrong results);
DEEPINFER_NEURON_PROFILE_RANKS keeps NEURON_RT_INSPECT on for the listed followers.
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>
A world all-reduce costs ~3.2 ms per graph that holds one, whatever it carries
(microbenchmark at 64 ranks: one all-reduce a graph 3.5 ms, ten chained in one
graph 3.6 ms; fp32 / bf16 and 1 / 6 rows alike), so the step pays per
collective-bearing graph, not per byte. The MoE's gate, experts, all-reduce
and shared expert now stay one graph at decode sizes (prompt chunks still
split). bench62 (512 in, TP=64, nextn=5): engine step 805 -> 740 ms.
Random-token prompts make greedy outputs chaotic, so a graph-boundary change
that moves bf16 rounding changes the generated tokens and the acceptance; the
step time is the comparable figure.
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>
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>
- 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>
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>
- 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>
- 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>
