- 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>
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>
