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>
English first, Chinese at the end in a collapsed <details> block. The PR
template lists the tests/neuron suite commands.
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>
- Engram: gather rows where the table lives (host by default, read via UVA
on GPU) and move them to the compute device.
- EPHandle: the topk_idx version guard tolerates inference-mode tensors on
the neuron device that do not report is_inference().
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pure-torch sparse attention followed FlashMLA's tests/ref.py (fp32 P).
The CUDA kernels (sm100 prefill phase1.cuh:238-253) and the official
reference round P = exp(s - running_max) to bf16 before the P@V GEMM in an
online softmax; do the same with 64-key blocks.
Localized with examples/neuron/probe_layer0.py (per-op reference vs DeepInfer
capture): layer 0 is now bit-exact up to two rounding points that are
DeepInfer's own GPU/Ascend behaviour (SwiGLU -> FP8 from fp32, bf16 routed
partial in the MoE combine). CPU end-to-end on the tiny checkpoint: 26/27
greedy tokens match the reference, the miss being a 0.0096-margin step the
reference flips against itself.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout,
pure-torch replacements for the official reference kernels, golden
generation, a teacher-forced DeepInfer runner (LLM offline API) and a
comparison tool. CPU mode: 25/27 greedy tokens match, worst logits
cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060.
- model_loader: check parameters against the platform device, not is_cuda.
- deep_gemm paged MQA logits: do not dereference block-table slots past the
context.
- bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would
serialize host compute).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- vendor/deep_ep: EPBuffer dispatch/combine with torch.distributed uneven
all_to_all (host-synced split sizes), EngramBuffer sharded lookup,
BucketBuffer all_gather (CP-only reduce paths stubbed); verified against
the DeepEP-Ascend references at world size 1/2/4.
- tilelang_ops/neuron sampling: random_sampling (asc two-level fixed-point
sampler, Philox), fused_sample (logprobs / do_sample), sampler metadata,
sum_smaller_probs, structured-output bitmask; DSpark/MTP/varlen kernels
raise NotImplementedError. All 42 asc *_tl factories now have a Neuron
counterpart.
- vendor/flash_attn: SDPA varlen stand-in for the ViT (mirrors npu_ops).
- transfer_to_neuron: keep `with torch.device(...)` working on patched
factories; fix manual_seed recursion in host mode.
- engram: gather the packed ue8m0 sf as bytes on Neuron.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Run DeepInfer on the native torch-neuronx device (PrivateUse1 `neuron`,
torch-neuronx 2.14, no torch-xla), following the Ascend port's structure.
- deep_cluster: Neuron discovery (one DeepInfer device = one logical
NeuronCore, NEURON_RT_VISIBLE_CORES, NUMA from neuron-ls).
- PlatformNeuron: takes the Ascend (is_npu) kernel layouts, adds is_neuron;
DEEPINFER_NEURON_DEVICE=cpu runs the same code on the host for
correctness work.
- neuron_ops/transfer_to_neuron: map "cuda" devices, torch.device and
torch.cuda.* onto neuron:0 (counterpart of torch_npu transfer_to_npu).
- C++ extension: -DWITH_NEURON host-only build; torch-allocated buffers and
tensor-view slot copies instead of raw device-pointer copies.
- Pure-torch drop-in vendor packages (Ascend API semantics): deep_gemm,
flash_mla, deep_select, tile_kernels, tilelang (constants only).
- tilelang_ops/neuron backend: KV cache / metadata / norm / compress kernels;
CP, ZOS and VL entry points raise NotImplementedError.
- Basic mode defaults on Neuron: eager, no DSpark/MTP drafting, no
zero-overhead scheduling; MoE avoids the GPU-only mega kernel; mHC
prenorm GEMM uses the Ascend single-split reduction.
Not yet included: deep_ep, the sampling kernels, end-to-end harness.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
