lightseekorg/kimi-k3-dspark

Kimi-K3 DSpark

Kimi-K3 DSpark is a speculative draft model for accelerating Kimi-K3. It was trained with TorchSpec, using online target-model inference and FSDP draft-model training. The architecture follows DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation.

The draft uses a five-layer MLA backbone, five target hidden-state taps, a low-rank Markov head, and a confidence head. It proposes up to seven draft tokens per speculative decoding step.

Training recipe

The training recipe was adapted from deepseek-ai/DeepSpec for online Kimi-K3 target inference and distributed TorchSpec training.

Due to limited compute resources, we trained on only approximately 150K selected and regenerated samples. We selected the prompts from lightseekorg/kimi-mtp-dataset and regenerated the responses with Kimi-K3. Regeneration is important for speculative decoding because the draft must learn the token-level distribution of the exact target model it accelerates.

The regenerated mix covers general instruction following, reasoning, coding and software engineering, tool use, Chinese, and multilingual prompts. Training used the Kimi-K3 renderer with assistant-only loss on sequences up to 32K tokens.

Setting	Value
Training epochs	5
Optimizer steps	6,891
Global batch size	128
Maximum sequence length	32,768
Learning rate	5e-5, cosine schedule
Warmup	4%
Loss objective	Decayed DSpark loss, gamma=4.0
Token loss weights	Cross-entropy=0.1, L1=0.9
Confidence-head loss weight	1.0
Draft layers	5
Target hidden-state layers	7, 31, 47, 63, 87
Speculative tokens	7

Compute cost

The end-to-end online training run used 40 NVIDIA GB200 GPUs: 8 GPUs for FSDP training and 32 GPUs for online Kimi-K3 inference, equivalent to approximately 1,104 GB200 GPU-hours.

Performance

The primary metric is acc_len, defined as 1 + accepted_draft_tokens / draft_steps. Higher is better. Evaluation used num_speculative_tokens=7 and Kimi-K3 thinking mode.

Benchmark	Samples	Greedy draft, t=0	Probabilistic draft, t=1.0, top_p=0.95
GSM8K	1,319	5.174	4.947
HumanEval	164	4.585	4.025
SPEED-Bench coding	80	3.986	3.745
SPEED-Bench RAG	80	3.373	3.265
SPEED-Bench multilingual	80	3.279	3.029
MATH-500	500	3.098	3.076
SPEED-Bench QA	80	2.895	2.762
MT-Bench	80	2.871	2.699
SPEED-Bench writing	80	2.555	2.426
AIME 2026	30	2.430	2.237

Requests were sent through /v1/chat/completions with:

{
  "thinking": true,
  "thinking_effort": "max"
}

Launch with TokenSpeed

This checkpoint uses native DSpark support in TokenSpeed. The auxiliary AttnRes stream must be enabled for this training contract.

export TOKENSPEED_DFLASH_AUX_STREAM=attn_res

tokenspeed serve moonshotai/Kimi-K3 \
  --served-model-name moonshotai/Kimi-K3 \
  --tensor-parallel-size 16 \
  --distributed-executor-backend ray \
  --trust-remote-code \
  --max-model-len 131072 \
  --gpu-memory-utilization 0.92 \
  --speculative-config '{"method":"dspark","model":"lightseekorg/kimi-k3-dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_MLA","draft_sample_method":"greedy","rejection_sample_method":"block"}' \
  --host 0.0.0.0 \
  --port 8000

Use "draft_sample_method":"greedy" for temperature-zero workloads. For sampled generation, launch a separate endpoint with "draft_sample_method":"probabilistic". The draft sampling method is a server-level setting.

Important serving notes

1. The collected target hidden states must be captured after the attention residual (attn_res) has been applied and before the following pre-normalization. Set TOKENSPEED_DFLASH_AUX_STREAM=attn_res; using another residual stream changes the draft model's input distribution.
2. This checkpoint's config declares fc_norm. Each collected hidden-state tap must pass through its corresponding fc_norm before the five taps are concatenated and passed to the linear context projection. Skipping this step is incompatible with the training contract.

TokenSpeed support for both requirements is implemented in lightseekorg/tokenspeed#1016.

Citation

@misc{cheng2026dspark,
  title={DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation},
  author={Xin Cheng and others},
  year={2026},
  eprint={2607.05147},
  archivePrefix={arXiv},
  primaryClass={cs.AI},
  url={https://arxiv.org/abs/2607.05147}
}

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