380e160b4b
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>
3.0 KiB
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.