mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00

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

5.1 KiB

Raw Blame History

DeepInfer on AWS Trainium (native torch-neuronx, no XLA)

Stack

torch 2.14 + torch-neuronx 2.14.3 (PrivateUse1 device neuron, eager + torch.compile(backend="neuron")), Neuron SDK 2.33 (runtime 2.35). No torch-xla, no libtorch-neuronx-lite.
Target: trn2.48xlarge (64 logical NeuronCores, LNC=2, 24 GB HBM per core), Neuron DLAMI.
Python env: a venv layered on the DLAMI's /opt/aws_neuronx_venv_torch_neuronx_2_14_0 (add its bin/ to PATH so neuronx-cc is found), plus requirements.txt minus tilelang / tile-kernels / torch.
Build: DEEPINFER_ACCELERATOR=neuron ./scripts/build.sh build (host-only C++ extension, -DWITH_NEURON).

Platform design

PlatformNeuron (deepinfer/platform.py): is_npu is true (Neuron follows the Ascend kernel layouts / code paths), is_neuron is true, tilelang_ops_backend = "neuron".
deepinfer/neuron_ops/transfer_to_neuron.py maps "cuda" devices and torch.cuda.* onto the target device. DEEPINFER_NEURON_DEVICE=neuron|cpu selects the target; cpu runs everything on the host.
deepinfer/neuron_ops/vendor/ is put at the front of sys.path (deepinfer/init.py) so import deep_gemm, flash_mla, deep_ep, deep_select, tile_kernels, tilelang resolve to the pure-torch packages there.
deepinfer/model_executor/layers/tilelang_ops/neuron/ is the kernel backend: every xxx_tl factory that asc/ provides, same name and signature, returning a callable with the same call signature and the same in-place output semantics.

Scope: basic adaptation only

Optional optimizations stay OFF on Neuron (defaults set in deepinfer/init.py): CUDA graphs (eager), DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0), zero-overhead scheduling, MTP varlen, EPLB, CP>1, YOCO bounded replay, batch-invariant mode, GPU-only mega kernels, VL/vision. Kernels reachable only from those features: export the name with the right signature but raise NotImplementedError("<name>: not supported on Neuron (basic mode)"). Fully implement what the model architecture needs (MLA attention, DSA indexer/top-k, compressor, mHC, Engram, MoE with FP8/FP4 dequant, norms/rope, KV cache store/load, basic sampling).

Rules for kernel code (vendor packages and tilelang_ops/neuron)

Pure torch, device-agnostic: allocate on the input tensors' device; never hardcode cuda/npu/neuron. Must run with tensors on CPU (correctness now) and on neuron (later).
Semantics follow the Ascend variant (DeepGEMM-Ascend, DeepEP-Ascend, tilelang_ops/asc/, the PLATFORM.is_npu branches); where Ascend has no variant, follow the CUDA one.
Output semantics must match exactly: if a kernel writes into caller-provided tensors, write in place (out.copy_(...), indexed assignment); keep the caller-visible shape/dtype/stride contract.
Prefer static shapes and no host syncs (.item(), .tolist(), nonzero, boolean-mask indexing) in per-step model code — the decode step will later be compiled to one static graph per batch bucket. Where a data-dependent shape is unavoidable, isolate it and leave a # NEURON: host sync comment.
Neuron dtype notes: prefer int32 over int64 for indices on device; fp8 is torch.float8_e4m3fn; FP4 (e2m1) is packed two per byte — unpack with bit ops; ue8m0 scales are exponent bytes (value = 2**(b-127)).
Correctness over speed. Match the reference numerics (fp32 accumulation where the kernel accumulates in fp32).
Each file has a short module docstring saying which upstream API it mirrors.
float8 on the neuron device: torch-neuronx 2.14 returns wrong data for torch.cat/stack of float8 tensors (e4m3fn and e8m0fnu) and for advanced indexing (t[idx_tensor, ...]) on non-contiguous float8 views (tests/neuron/device/test_device_primitives.py); the same ops on uint8 views of the same bytes are correct. Move 1-byte floats as bytes, for every concat / gather / scatter / index_put, contiguous or not: torch.cat([t.view(torch.uint8) for t in parts], dim).view(dtype), t.view(torch.uint8)[idx].view(dtype), dst.view(torch.uint8)[idx] = src.view(torch.uint8) (helpers: deep_ep/utils/indexing.py). This is bit-identical on every platform, so loader/model code can use it unconditionally (read_concat, merge_linear_weights). Basic slicing, dtype casts and elementwise math followed by an upcast are fine. To find offending sites without a NeuronCore, run with DEEPINFER_NEURON_FP8_GUARD=1 (also in CPU mode): deepinfer/neuron_ops/fp8_guard.py records the call site (file:line, dtype/shape/contiguity) of every float8 cat/stack/advanced-index/gather/scatter/index_put/where and prints a deduplicated summary at exit (=raise raises at the first one; DEEPINFER_NEURON_FP8_GUARD_LOG=<file> appends the summary to a file). The tiny e2e run and the CPU unit stages must report 0 sites.

Testing

CPU mode (same code, tensors on the host):

export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu

Compare against the upstream references (*_ref functions next to each tilelang_ops interface, tile_kernels.torch, and the reference implementations inside vendor/<Repo>/tests).

5.1 KiB Raw Blame History