coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 23:07:09 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/examples/neuron
History
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
..
golden Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
check_tiny_checkpoint.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
compare_golden.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
device_diff.py Neuron: tiny model on a NeuronCore matches the reference (issues #1, #2, #4) 2026-10-21 13:22:59 +08:00
make_tiny_checkpoint.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
probe_layer0.py Neuron: tiny model on a NeuronCore matches the reference (issues #1, #2, #4) 2026-10-21 13:22:59 +08:00
README.md Neuron flash_mla: round attention P to bf16 like the CUDA kernels 2026-10-20 20:48:36 +08:00
ref_kernels_torch.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
reference_golden.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00
run_deepinfer_tiny.py Neuron: EP > 1 with real neuron collectives (issue #6) 2026-10-21 11:46:59 +08:00
test_ref_kernels_torch.py Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +08:00

README.md

DeepInfer on AWS Trainium — end-to-end check with a tiny checkpoint

Design and conventions of the port: deepinfer/neuron_ops/NEURON_PORT.md.

The ground truth is DeepSeek's official reference implementation (inference/ in the deepseek-ai/DeepSeek-V4.1-Flash-1017 HF repo) running on CPU, with its TileLang kernels replaced by the pure-torch versions in ref_kernels_torch.py.

Script Purpose
make_tiny_checkpoint.py Random tiny V4.1 checkpoint in the real HF layout (FP8 e4m3 + ue8m0 32x32 scales, FP4 e2m1 experts with 1x32 scales, engram tables)
check_tiny_checkpoint.py DeepInfer loader coverage / shape / dtype checks on the meta device
ref_kernels_torch.py, test_ref_kernels_torch.py Torch replacements for the reference kernel.py
reference_golden.py Reference prefill + greedy decode; saves logits, tokens and per-layer hidden states
run_deepinfer_tiny.py Same prompts through the DeepInfer engine (LLM offline API), teacher-forced
compare_golden.py Candidate vs golden: argmax match, logits rel / cos / KL, per-layer hidden error

Tiny config: 10 layers + 3 nextn, hidden 1024, 8 heads x 512 (MLA, rope 64), 16 experts / top-4, indexer 4 x 128 with top-64, engram on layers 1 and 4, real tokenizer (vocab 129280).

export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu   # or =neuron for the NeuronCore
python examples/neuron/run_deepinfer_tiny.py --threads 32 --hidden --compare

Results so far (CPU mode, basic mode: eager, no DSpark/MTP, EP=DP=CP=1)

Teacher-forced, prefill + 8 decode steps per prompt; "floor" is the reference compared with itself under a different CPU thread count (FP8 activation quantization amplifies sub-ulp differences).

prompt (tokens) argmax match max logits rel (floor) min cos (floor)
short (11) 9/9 0.054 (0.032) 0.9985 (0.9995)
medium (172, wraps the 128 SWA window) 8/9 * 0.041 (0.039) 0.9992 (0.9993)
long (715, candidate-block filter active) 9/9 0.059 (0.060) 0.9983 (0.9982)

* medium step 6 has a top-1/top-2 margin of 0.0096; the reference flips it against itself too.

probe_layer0.py compares ~35 named ops per layer between the reference and DeepInfer. Layer 0 is bit-exact once three rounding points are aligned:

  • attention P rounded to bf16 before P@V (online softmax, 64-key blocks) — the CUDA FlashMLA kernels and the reference do this; the Neuron flash_mla shim now does too (was fp32, a port bug);
  • SwiGLU output quantized to FP8 straight from fp32 (the reference rounds to bf16 first) and
  • the MoE combine rounding the routed sum to bf16 before adding the shared expert (the reference adds in fp32) — both are DeepInfer's own GPU/Ascend behaviour, kept as is.

The remaining logits difference comes from those two, amplified by FP8 activation quantization layer by layer. Per-step / per-layer numbers before the attention fix: golden/deepinfer_cpu_vs_golden.json.