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

Design and conventions of the port: [deepinfer/neuron_ops/NEURON_PORT.md](../../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).

```bash
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`.