whn09 8a5209ee65 Neuron basic mode — part 3: end-to-end check vs the official reference
- 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>
2026-10-20 20:20:02 +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
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
README.md Neuron basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +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 basic mode — part 3: end-to-end check vs the official reference 2026-10-20 20:20:02 +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

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.055 (0.032) 0.9985 (0.9995)
medium (172, wraps the 128 SWA window) 8/9 0.049 (0.039) 0.9988 (0.9993)
long (715, candidate-block filter active) 8/9 0.056 (0.060) 0.9985 (0.9982)

Per-step / per-layer numbers: golden/deepinfer_cpu_vs_golden.json.