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, trained with 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.

It is the final checkpoint of a completed one-epoch run, continued from 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: 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, 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	160	2.834
Chinese	CEval	1346	2.486
Multilingual	SPEED-Bench multilingual	80	2.986
Math	GSM8K	1319	3.571
Math	MATH500	500	3.151
Math	AIME 2026	30	2.262
Code	HumanEval	164	3.326
Code	SPEED-Bench coding	89	3.101
RAG	SPEED-Bench rag	91	3.050
QA	SPEED-Bench qa	80	2.739
Writing	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 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, or vLLM >= 0.20.0

Launch Server (TokenSpeed)

This is the configuration the numbers above were measured on.

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)

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)

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

@misc{torchspec2026,
  title={TorchSpec: An Online Speculative Decoding Training Framework},
  url={https://github.com/lightseekorg/TorchSpec},
  year={2026}
}