73 Commits

Author SHA1 Message Date
whn09
0a7d490cdb neuron: compiled decode verification behind DEEPINFER_NEURON_COMPILE_VERIFY (default off)
Sampler._verification runs as one compiled region for greedy batches without structured output: the native
sampling structs enter as Python mirrors (neuron_ops/sampling_mirror.py), the flags are read once per batch
eagerly (prime_host_flags), and greedy mirrors carry no seeds / positions. Untested on device yet.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 21:33:34 +08:00
whn09
b18af0027f neuron: sampling flags per request batch, NUMA-local CPU binding, sync/op-site tracing
- fused_sample: _host_flags/_greedy_verify flags persist while the batch's requests are unchanged
  (each device read drained the queue and waited for the whole forward). Tiny golden 9/9; full model
  step time unchanged (660 ms), driver dflash launch 71 -> 10.5 ms host.
- bind_cpu on Neuron: the device's NUMA node (DEEPINFER_NEURON_CPU_BIND=node|core|legacy); the generic
  split gave each rank one CPU and ranks 0-15 the wrong node. Neutral on TPOT (228.2 ms).
- profiled followers exit through sys.exit so the runtime writes their trace.
- op_census: trace_syncs / census_sites (DEEPINFER_NEURON_TRACE_SYNCS=1|2|3).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 19:04:53 +08:00
whn09
b6d8df7484 neuron: diagnostics for MoE collective skew (expert flag force, per-follower runtime profiling)
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>
2026-10-26 15:14:19 +08:00
whn09
d7d2f36ec1 neuron: float-only fp8 codec in the compile backend (no GpSimd int ops / remainder)
Exact for all 256 e4m3fn codes and RNE encode on host; tiny golden 9/9.
vllm-sonnet 512: TPOT 228.6 -> 223.3 ms, TTFT 4349 -> 4117 ms, tokens/step unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 12:41:42 +08:00
whn09
e7cb740b97 Neuron attention: decode fp4 KV records in float arithmetic (Vector engine), not integer bit ops (GpSimd)
The compressed-KV scopes are fp4 records; their decode was nibble masks and
shifts plus the integer e2m1 / e4m3 decoders on GpSimd (three such chains were
the largest cost of the ratio-2 attention graph, device profile bench72). Now:
nibbles from floor(b / 16), e2m1 from a piecewise-linear map of the magnitude
code, scales through the hardware fp8 conversion. Bit-identical to the integer
decode for every byte (host check). bench73 (vllm-sonnet 512, TP=64, nextn=5):
TPOT 243.5 -> 228.6 ms, TTFT 4.84 -> 4.35 s, step 699 -> 666 ms, tokens per
step identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 08:18:39 +08:00
whn09
94fa5c92c0 Neuron table attention: decode a whole cache in place instead of an arange gather
A scope whose blocks the step does not compact (the SWA cache) was gathered with
an arange through gather_rows: integer index math and an element-indexed GpSimd
gather over every slot. A contiguous cache in slot order is the cache itself,
so it is decoded in place. bench71 (vllm-sonnet 512, TP=64, nextn=5): TPOT
248.4 -> 243.5 ms, step 711 -> 699 ms, tokens per step identical; tiny cp=2
nextn=5 on device 9/9.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 07:17:36 +08:00
whn09
3c38fd5829 Neuron MoE: skip unrouted experts inside the fp4 kernels at decode sizes (DEEPINFER_NEURON_GATED_EXPERTS, default on)
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>
2026-10-26 06:36:48 +08:00
whn09
7e032422d2 Neuron attention: decode fp8 KV records with the hardware fp8 conversion (DEEPINFER_NEURON_FAST_E4M3, default on)
The record decode was a dozen integer GpSimd ops per element. Viewing the bytes
as float8_e4m3fn (read as trn2 e4m3 under the existing compile flag; DeepInfer
quantizes to 224 on Neuron, where the encodings agree) is one conversion.
Outputs identical (tiny device golden bit-identical; full-model tokens per
step identical), bench68 (vllm-sonnet 512): TPOT 262.0 -> 258.0 ms, TTFT
4.98 -> 4.90 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 05:00:07 +08:00
whn09
1dce2c4995 Neuron bench: vllm-neuron#71's sonnet prompts, prompt/output dumps; DSpark draft ground truth
- bench_neuron.py --dataset vllm-sonnet: vLLM's SonnetDataset as #71's bench
  builds it for the completions backend ("Pick as many lines..." + random poem
  lines, chat-formatted; --chat_tokenizer for a template the checkpoint lacks).
  --dataset sonnet: a contiguous text window. --dump: prompt/output token ids.
- reference_dspark.py: the official inference/model.py forward_spec on the host
  CPU (tiny checkpoint, teacher-forced on reference_golden.py tokens): each
  decode step's 5 drafts and their logits.
- run_deepinfer_tiny.py RUN_TINY_DRAFTS=1 prints every draft block.

DeepInfer's drafts match the reference on host eager, on device and with
attention TP: every divergence is a near tie (top-1/top-2 gap 0.001-0.055).
With #71's prompt format the acceptance is #71's: 2.8 tokens a step at 512
(first position 65-79%); a raw text window gave 1.4 (continuations not in the
context), random token ids 3.85 (repetitive outputs).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-26 03:16:54 +08:00
whn09
39f8aa6eed Neuron DSpark: 6-row steps run the MoE as one region graph (5 graphs a layer)
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>
2026-10-26 00:01:52 +08:00
whn09
5e36f90015 Neuron table attention: decode only the blocks a decode step's requests reference
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>
2026-10-25 17:58:11 +08:00
whn09
0016315f80 Neuron: eager embedding/gather without dtype casts or host scalar operands
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>
2026-10-25 06:55:43 +08:00
whn09
ee9e3aaed9 Neuron DSpark: draft sampling as one region graph for greedy rows; token embedding on device by default
- 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>
2026-10-25 06:06:50 +08:00
whn09
e6c7d7867a Neuron: sampling flags read once a step; runtime profiling on the driver only; one-lookup embedding
- 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>
2026-10-25 03:52:47 +08:00
whn09
976f9fe4d2 Neuron DSpark: 6-row steps keep only coarse region splits (8 graphs a layer instead of 12)
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>
2026-10-25 00:31:45 +08:00
whn09
afd49fe8ce Neuron DSpark: pad multi-row decode to a power of two before the attention path; 6-row steps take the table path
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>
2026-10-24 22:31:05 +08:00
whn09
e0510e92ca Neuron DSpark: bench --nextn step count fix; split-row threshold knob; table attention stays prompt-only
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>
2026-10-24 21:14:58 +08:00
whn09
fdbb81f14f Neuron DSpark (#15): draft layers and main-cache writes as region graphs, decode rows padded to a power of two
- 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>
2026-10-24 13:41:17 +08:00
whn09
b2d94e16b2 Neuron DSpark (#15): greedy verification and MTP ops, CP followers run the drafts
- 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>
2026-10-24 11:31:59 +08:00
whn09
f67fab6bdd Neuron MoE: two-level all-reduce (device, then across devices) (#27)
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>
2026-10-24 11:05:25 +08:00
whn09
5e89a12997 Neuron: whole-layer region graphs by default (TPOT 447.8 -> 430.3 ms at TP=64)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-24 10:15:43 +08:00
whn09
3a8ef041db Neuron: optional whole-layer region graph (DEEPINFER_NEURON_REGION_LAYER)
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>
2026-10-24 09:51:01 +08:00
whn09
e5663a0a1d Neuron: attention TP over lanes of 8 consecutive ranks (DEEPINFER_NEURON_ATTN_TP) (#28)
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>
2026-10-24 09:16:01 +08:00
whn09
2310c2db8a docs/neuron/README.md: running DeepSeek-V4.1-Flash on trn2, performance
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-24 08:31:57 +08:00
whn09
57d373e25b Neuron chunk attention: decode the cache once, gathered_attn_idx (the #71 way) (#26)
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>
2026-10-24 07:29:00 +08:00
whn09
2836e365ea Neuron engram shared tables: one file for the whole group (#28)
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>
2026-10-24 07:08:37 +08:00
whn09
d44cf253c5 offline_inference_neuron.py --cp
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-24 00:37:23 +08:00
whn09
3ee74e8a4e Neuron CP followers: keep the EPLB step when skipping the sampler (#28)
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>
2026-10-24 00:04:02 +08:00
whn09
ad6b47a141 Neuron: head-sharded attention over the CP group (TP decode) (#28)
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>
2026-10-23 22:53:32 +08:00
whn09
f9cbf6f2e5 Neuron: CP group as a replicated-decode group; MoE all-reduce mode (#28)
- 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>
2026-10-23 22:46:04 +08:00
whn09
d188d1b3d1 Neuron: split region graphs for prefill chunks only (#26)
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>
2026-10-23 21:39:56 +08:00
whn09
32e08d78cd bench_neuron: one untimed request a cell (context-bucket graphs compile lazily)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 21:16:45 +08:00
whn09
50483a06df Neuron: context-length buckets for dense compressed-KV scopes and indexer logits (#29)
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>
2026-10-23 20:08:17 +08:00
whn09
390ff603e0 Neuron chunk attention: NKI sparse kernel for top-k/SWA-only scopes (#26)
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>
2026-10-23 18:20:48 +08:00
whn09
279334f330 Neuron: capacity MoE off by default (measured slower) (#26)
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>
2026-10-23 17:48:39 +08:00
whn09
ad23d2d870 Neuron MoE capacity path: all-reduce the overflow flag; bench prompts differ per rank (#26)
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>
2026-10-23 17:06:47 +08:00
whn09
7d95b7b6d6 Neuron: count MoE capacity overflows under DEEPINFER_NEURON_REGION_TIMING
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 16:23:28 +08:00
whn09
d8f7a2e898 Neuron: PR #147 OpenAI server (+ /v1/completions, Neuron worker setup), offline bench, MoE fixes (#26, #27, #29)
- 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>
2026-10-23 16:23:28 +08:00
whn09
7d897aedb2 Neuron: skip the unused candidate-block source when the filter is vacuous (#26)
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>
2026-10-23 16:23:28 +08:00
王鹤男
5dfa098dca
Merge branch 'deepseek-ai:main' into neuron 2026-10-23 16:21:31 +08:00
whn09
fb9dea85cb Neuron sampler: skip the top-p/top-k bisection when no row filters (#27)
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>
2026-10-23 12:05:16 +08:00
whn09
e846322c0c Neuron: block-quant GEMMs on both physical cores (grid 2); sync experiment knob (#27)
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>
2026-10-23 11:08:05 +08:00
whn09
f07d0a7c6f Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27)
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>
2026-10-23 09:42:08 +08:00
whn09
563c181ecc Neuron: integer mod without torch.remainder in the sampler; sampler cProfile (#27)
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>
2026-10-23 08:10:16 +08:00
whn09
42e05844d1 Neuron: nkilib rotational_topk for the sampler; single-host shared engram tables (#27)
- 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>
2026-10-23 07:23:32 +08:00
whn09
a263b60514 Neuron: profile a prefill chunk too under DEEPINFER_NEURON_REGION_TIMING
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 06:38:31 +08:00
whn09
32abff6662 Neuron: skip the vacuous candidate-block filter; mask padding rows in static MoE (#26)
- 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>
2026-10-23 05:25:28 +08:00
whn09
0cae424dc4 Neuron: one KV gather per chunk for dense attention; no layer_idx guards (#26)
- 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>
2026-10-23 04:32:20 +08:00
whn09
4bc8ed3eaf Neuron: separate state regions, chunk-wide indexer context (issue #26)
- 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>
2026-10-23 03:46:56 +08:00
whn09
3629d113fb Neuron: prompt chunks through the decode region graphs (issue #26)
- 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>
2026-10-23 03:19:56 +08:00