mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
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>
50 lines
3.0 KiB
Markdown
50 lines
3.0 KiB
Markdown
# 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`.
|