- 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>
Uneven all_to_all on the neuron process group is unreliable at scale (EP=64:
"Invalid NEFF" in prefill; test_ep at 64 ranks hung after 9/18 modes).
EPBuffer now moves tokens without all_to_all:
- dispatch: one all_gather of x / sf / weights / topk_idx packed as uint8
rows (row count = global max via a 0-dim MAX all_reduce, or the static
num_max_tokens_per_rank with DEEPINFER_NEURON_EP_AG_ROWS=cap); each rank
selects the rows routed to its local experts in the same (source rank,
source token) order as before, so layout / metadata / cached dispatch are
unchanged.
- combine: per-token fp32 partials over local experts (bf16-rounded as
before) + weights, one fp32 SUM reduce_scatter (all_reduce fallback with
DEEPINFER_NEURON_EP_REDUCE=ar), chunked under
DEEPINFER_NEURON_EP_CHUNK_BYTES (64 MiB).
- The a2a path stays behind DEEPINFER_NEURON_EP_ALGO=a2a.
test_ep passes on the neuron group at 2/4/8/16/32/64 ranks (combine
bit-identical to the reference in practice); new test_collectives.py found
that int32 all_gather with full-range values is wrong on a 16-rank neuron
group for >= 1 MiB outputs (uint8 / fp32 fine) -- payloads move as bytes.
Tiny model EP=8 on NeuronCores: 25/27 argmax vs the reference, same as EP=1.
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>
- deep_ep comm: torch-neuronx picks the all_to_all lowering per rank from
that rank's own splits (even / NKI a2a-v / 1-D-only padding fallback), so
ranks could diverge. All-gather the split matrix first and choose one path
for all ranks: even a2a, NKI a2a-v for groups of 4/8, otherwise a padded
even a2a (2 ranks, 16 ranks where a2a-v fails on runtime 2.35).
- PlatformNeuron: LOCAL_RANK=0 / LOCAL_WORLD_SIZE=1 so every one-core
process is neuron:0 of a 1-device local world (otherwise rank > 0 died in
_lazy_init at EP/DP > 1).
- tests/neuron/deep_ep: references on host copies with the gloo group,
buffers on the device; ep/engram/bucket pass at world size 2/4/8/16 on the
neuron process group, and 1/2/4 on gloo.
- run_deepinfer_tiny.py --ep N: one process per rank / NeuronCore; tiny
model EP=2 runs end to end (CPU: within the noise floor).
- run_all.sh: per-core busy check, bucket + world size 16 device stages.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
torch.cat of float8 tensors (e4m3 and e8m0) returns wrong data on the
neuron device (torch-neuronx 2.14). The HF loader concatenates w1/w3 (and
their UE8M0 scales) into the fused gate_up weights, so on the NeuronCore
every routed and shared expert gate_up scale was zero and the MoE output
vanished (issue #1). Concatenate 1-byte float dtypes through uint8 views,
which is bit-identical on every platform.
tests/neuron/device: create the neuron process group before the runtime
initializes (torch-neuronx requirement) and add the float8 cat checks.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
