5 Commits

Author SHA1 Message Date
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
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
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
Jiashi Li
d585e6542c initial commit 2026-10-20 12:09:42 +08:00