Device-memory fixes found bringing up the released checkpoint (16.45 GiB of
weights per 24 GB core):
- Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the
all-gather of the replicated SF table (~3 GB per layer, allocated on the
device -> "Failed to allocate resource").
- gather_tl: index_select / advanced indexing of an int8 table allocates
several times the table on the device (+4.9 GiB for the 1.3 GB embedding);
gather through F.embedding on an int32 / fp16 view instead (bit-exact, no
extra memory).
- deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements
and one expert at a time instead of whole projections / all local experts
in fp32 (OOM on a fragmented pool); bit-identical results.
- transfer_to_neuron: import torch.utils._device explicitly (missing when
torch_neuronx is not auto-loaded).
examples/neuron/offline_inference_neuron.py: offline inference in basic mode
(text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for
the candidate-block filter), one process per NeuronCore.
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>
- 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>
