whn09 8935f80230 Neuron: tiny model on a NeuronCore matches the reference (issues #1, #2, #4)
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>
2026-10-21 13:22:59 +08:00
..
eplb initial commit 2026-10-20 12:09:42 +08:00
input_metadata initial commit 2026-10-20 12:09:42 +08:00
layers Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
models Neuron: tiny model on a NeuronCore matches the reference (issues #1, #2, #4) 2026-10-21 13:22:59 +08:00
parallel_utils initial commit 2026-10-20 12:09:42 +08:00
__init__.py initial commit 2026-10-20 12:09:42 +08:00
layer_utils.py initial commit 2026-10-20 12:09:42 +08:00
model_loader.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
model_metadata_utils.py initial commit 2026-10-20 12:09:42 +08:00
model_output.py initial commit 2026-10-20 12:09:42 +08:00
sampling_inputs.py initial commit 2026-10-20 12:09:42 +08:00
utils.py initial commit 2026-10-20 12:09:42 +08:00