| license: other | |
| base_model: moonshotai/Kimi-K3 | |
| library_name: transformers | |
| tags: | |
| - text-generation | |
| - speculative-decoding | |
| - eagle3 | |
| - kimi-k3 | |
| - torchspec | |
| # kimi-k3-eagle3.1-mla | |
| ## Model Overview | |
| kimi-k3-eagle3.1-mla is an EAGLE-3.1 MTP draft model with MLA (Multi-Latent Attention) for accelerating inference of [Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3), trained with [TorchSpec](https://github.com/lightseekorg/TorchSpec) — an online speculative decoding training framework that runs FSDP training and inference concurrently. If you find this draft model useful, please give our project TorchSpec a star on [GitHub](https://github.com/lightseekorg/TorchSpec). | |
| It is the final checkpoint of a completed one-epoch run, continued from [`lightseekorg/kimi-k3-eagle3-mla`](https://huggingface.co/lightseekorg/kimi-k3-eagle3-mla). | |
| ### Why an MLA (Multi-Latent Attention) Draft Model | |
| Compared with an MHA draft model, the MLA variant is a better fit for Kimi-K3 deployment: | |
| - Uses less KV cache, which reduces serving memory pressure. | |
| - Matches Kimi-K3's MLA architecture, so it fits more naturally into the inference engine's KV-cache handling under different serving scenarios such as PD-Disaggregation. | |
| ### What EAGLE-3.1 Adds Over EAGLE-3 | |
| EAGLE-3.1 applies the normalization scheme introduced in [Attention Drift: What Autoregressive Speculative Decoding Models Learn](https://arxiv.org/abs/2605.09992): a post-norm on the drafter's hidden states, and a per-hidden-state RMSNorm applied to each captured target hidden state before fusion. | |
| Both address the same failure mode. An EAGLE-style drafter's hidden state magnitude grows monotonically with the depth of the speculation chain, so the drafter starts behaving like another stack of target layers rather than a standalone predictor. As magnitude grows, attention drifts off the prompt and onto the tokens the drafter has just generated. Constraining that growth gives three practical improvements: | |
| - **Higher acceptance.** More tokens accepted per verification step across ordinary workloads — chat, math, and code. | |
| - **Out-of-distribution resilience.** The gain is largest exactly where EAGLE-3 drafters are weakest. The paper reports up to 2x acceptance under template perturbation, where a drafter trained on one chat template meets another. | |
| - **Slower acceptance decay.** Acceptance falls off far more gently as the draft chain deepens and as context grows, rather than collapsing at depth. This is what makes long-context serving hold up. | |
| ## Performance | |
| ### Acceptance Length | |
| The primary metric is accept_length — the average number of tokens accepted per speculation step at `topk=1, num_steps=3, num_draft_tokens=4`. Higher is better; the ceiling for this configuration is 4. | |
| Measured against a BF16 Kimi-K3 target on GB200 (TP=16, 4 nodes) served with [TokenSpeed](https://github.com/lightseekorg/tokenspeed), at `temperature=1.0, top_p=0.95`, reasoning effort max. Acceptance is attributed per request from the scheduler log, so N is the number of requests measured. | |
| | Category | Benchmark | N | Acc Len | | |
| | --- | --- | --- | --- | | |
| | Dialogue | [MTBench](https://github.com/lm-sys/FastChat/tree/main/fastchat/llm_judge) | 160 | 2.834 | | |
| | Chinese | [CEval](https://huggingface.co/datasets/ceval/ceval-exam) | 1346 | 2.486 | | |
| | Multilingual | [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) multilingual | 80 | 2.986 | | |
| | Math | [GSM8K](https://huggingface.co/datasets/openai/gsm8k) | 1319 | 3.571 | | |
| | Math | [MATH500](https://huggingface.co/datasets/HuggingFaceH4/MATH-500) | 500 | 3.151 | | |
| | Math | [AIME 2026](https://huggingface.co/datasets/MathArena/aime_2026) | 30 | 2.262 | | |
| | Code | [HumanEval](https://huggingface.co/datasets/openai/openai_humaneval) | 164 | 3.326 | | |
| | Code | [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) coding | 89 | 3.101 | | |
| | RAG | [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) rag | 91 | 3.050 | | |
| | QA | [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) qa | 80 | 2.739 | | |
| | Writing | [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) writing | 84 | 2.447 | | |
| The acceptance table published for `kimi-k3-eagle3-mla` was measured on an NVFP4 target on GB300 at TP=8. Target precision, hardware and TP degree all move accept_length, so the two tables should not be read as a controlled A/B. | |
| ### Throughput | |
| Speculation on versus off, same build, same hardware, same allocation, run uncontended on a mixed multi-turn workload with `max_tokens=32768`. Output token throughput in tokens/s. | |
| | Concurrency | EAGLE-3.1 tok/s | No speculation tok/s | Speedup | ITL (ms) spec / no-spec | | |
| | --- | --- | --- | --- | --- | | |
| | 1 | 206.6 | 96.3 | 2.15x | 4.62 / 10.08 | | |
| | 2 | 333.7 | 174.7 | 1.91x | 5.34 / 10.66 | | |
| | 4 | 566.6 | 310.2 | 1.83x | 6.36 / 11.80 | | |
| | 16 | 1263.5 | 838.2 | 1.51x | 11.65 / 16.99 | | |
| Measured accept_length on this workload was 2.66 and was flat across all four concurrency levels. As with any speculative decoder the gain is largest at low concurrency, where decode is memory-bandwidth bound and there is spare compute available to verify draft tokens. | |
| ## Training | |
| Trained with [TorchSpec](https://github.com/lightseekorg/TorchSpec) using **Anchored EAGLE-3**, which supervises the drafter from anchor positions sampled across the sequence instead of only from the tail. This exposes the drafter to a much wider spread of context lengths and chain depths per sample, and is a large part of why acceptance holds up at long context. | |
| | | | | |
| |---|---| | |
| | Target | `moonshotai/Kimi-K3` (mxfp4, 93 layers) | | |
| | Aux layers | 2, 46, 90 (one-based) | | |
| | Objective | LK^lambda (`lk_eta=3.0`); forward KL for the first 1000 steps | | |
| | Anchored EAGLE-3 | 512 anchors/sample, `max_gap=0` | | |
| | Sequence length | 46000 | | |
| | Optimizer | AdamW, peak LR 2e-5, cosine to 0, 1% warmup | | |
| | Epochs | 1 (218,109 steps, 51.0 h on 20x GB300) | | |
| | Draft vocab | 163840 (unpruned) | | |
| --- | |
| ## Quick Start | |
| ### Requirements | |
| - NVIDIA GPU with CUDA 12.0+ | |
| - [TokenSpeed](https://github.com/lightseekorg/tokenspeed), or [vLLM](https://github.com/vllm-project/vllm) >= 0.20.0 | |
| ### Launch Server (TokenSpeed) | |
| This is the configuration the numbers above were measured on. | |
| ```bash | |
| python3 -m tokenspeed.cli serve /models/Kimi-K3 \ | |
| --served-model-name moonshotai/Kimi-K3 \ | |
| --attn-tp-size 16 --moe-tp-size 16 --mm-encoder-tp-mode data \ | |
| --max-model-len 80000 --max-num-seqs 16 \ | |
| --gpu-memory-utilization 0.92 --trust-remote-code \ | |
| --attention-backend tokenspeed_mla --kda-backend cutedsl_kda \ | |
| --moe-backend flashinfer_trtllm --kv-cache-dtype fp8_e4m3 \ | |
| --speculative-algorithm EAGLE3 \ | |
| --speculative-draft-model-path /models/eagle3_1 \ | |
| --speculative-num-steps 3 --speculative-eagle-topk 1 \ | |
| --speculative-num-draft-tokens 4 \ | |
| --speculative-draft-model-quantization unquant \ | |
| --sampling-backend flashinfer \ | |
| --drafter-attention-backend tokenspeed_mla \ | |
| --max-cudagraph-capture-size 16 --cudagraph-capture-sizes 1 2 4 8 16 \ | |
| --host 0.0.0.0 --port 8000 | |
| ``` | |
| `/models/Kimi-K3` is the target checkpoint and `/models/eagle3_1` is this draft; point both at your local paths. | |
| `--sampling-backend flashinfer` keeps rejection sampling faithful to the requested `temperature` / `top_p`, which is required to reproduce the acceptance numbers above. | |
| ### Launch Server (vLLM) | |
| ```bash | |
| vllm serve moonshotai/Kimi-K3 \ | |
| --tensor-parallel-size 16 \ | |
| --speculative-config '{"model": "lightseekorg/kimi-k3-eagle3.1-mla", "method": "eagle3", "num_speculative_tokens": 3}' \ | |
| --trust-remote-code | |
| ``` | |
| ### Launch Server (SGLang) | |
| ```bash | |
| sglang serve \ | |
| --model-path moonshotai/Kimi-K3 \ | |
| --tp 16 \ | |
| --trust-remote-code \ | |
| --reasoning-parser kimi_k2 \ | |
| --tool-call-parser kimi_k2 \ | |
| --speculative-algorithm EAGLE3 \ | |
| --speculative-num-steps 3 \ | |
| --speculative-eagle-topk 1 \ | |
| --speculative-num-draft-tokens 4 \ | |
| --speculative-draft-model-path lightseekorg/kimi-k3-eagle3.1-mla | |
| ``` | |
| ## Serving Notes | |
| **Layer indexing.** `eagle_aux_hidden_state_layer_ids` is `[2, 46, 90]`, the one-based convention that both vLLM and TokenSpeed expect. An SGLang-native loader wants `[1, 45, 89]` instead. Getting this wrong does not error — it silently costs roughly 0.35 accept_length. | |
| **Aux capture mode.** Trained against vLLM's `prefix_only` aux capture (`VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=0`). Builds that default that flag on tap the AttnRes mixture instead, which is a different feature stream. Serve with `VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=0` to match training. | |
| ## Citation | |
| ```bibtex | |
| @misc{torchspec2026, | |
| title={TorchSpec: An Online Speculative Decoding Training Framework}, | |
| url={https://github.com/lightseekorg/TorchSpec}, | |
| year={2026} | |
| } | |
| ``` | |
Xet Storage Details
- Size:
- 8.75 kB
- Xet hash:
- e32239259931170051ed668e86b75719aba5ae25454d3353957318a1d55266e6
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.