mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
- 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> |
||
|---|---|---|
| .. | ||
| golden | ||
| check_tiny_checkpoint.py | ||
| compare_golden.py | ||
| make_tiny_checkpoint.py | ||
| README.md | ||
| ref_kernels_torch.py | ||
| reference_golden.py | ||
| run_deepinfer_tiny.py | ||
| test_ref_kernels_torch.py | ||
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.