Inferact/Kimi-K3-DSpark

Model Overview

Inferact/Kimi-K3-DSpark is an MLA-native DSpark draft model that accelerates inference of Kimi-K3 on vLLM, which serves it through its native dspark speculative method.

The draft is trained on target hidden states extracted from vLLM itself — the same engine that serves it — so it learns from exactly the distribution it will later draft against. TorchSpec provides the training loop, streaming those hidden states from live vLLM target inference into concurrent FSDP draft training.

DSpark = a block-diffusion backbone (5 dense layers that consume 5 target auxiliary hidden states and draft a block of 7 tokens in a single parallel pass), plus a low-rank sequential Markov head (rank=256) that supplies the intra-block token dependency the parallel backbone misses, plus a confidence head.

Why MLA: the draft mirrors Kimi-K3's own MLA attention, so draft and target share a single KV layout (a compact 576-element latent per token). The draft's pages unify with the target's cache, letting it reuse the same KV management and PD-disaggregated serving paths with no separate page format.

---

Performance

Measured end-to-end against Kimi-K3 at tensor-parallel-size=8 (2 × GB300), num_speculative_tokens=7. Sampling matches the production serving setting (temperature=1.0 / top_p=0.95), so acceptance reflects real serving conditions.

Acceptance tracks how predictable the output is: it is strongest on deterministic, structured text such as code, and weaker on high-entropy generation such as creative writing. Confidence-based scheduling — using the draft's confidence head to prune low-confidence drafted tokens before they are sent for verification — is planned future work aimed at lifting the harder categories.

Benchmarks: MT-Bench (80 dialogue prompts) and subset of NVIDIA SPEED-Bench — its throughput_16k low-entropy split (302 prompts at 10k-token input) plus five qualitative categories.

vLLM supports block verification (rejection_sample_method=block) and probabilistic draft sampling (draft_sample_method=probabilistic) — verifying the drafted block as a unit, and sampling the draft from its own distribution rather than its argmax. Both are enabled in the numbers above and push acceptance higher at temperature=1.0.

---

Training

Data — all responses regenerated on-policy by Kimi-K3 itself, so the draft learns the target's own reasoning traces and chat formatting. Prompts come from public datasets:

• lightseekorg/kimi-mtp-dataset — general instruction prompts
• nvidia/OpenCodeInstruct — single-turn coding
• A nine-category prompt pool from the NVIDIA Nemotron SFT/RL collections and CohereLabs/aya_dataset, spanning chat, code, multilingual, RAG/QA, math, structured output and safety. Evaluation-set prompts are excluded.

Method: DSpark with block_size=7, trained on a combined CE + L1 distribution-distillation objective against the target's post-final-norm hidden state, bf16. The draft consumes Kimi-K3 auxiliary hidden states from target layers (2, 23, 47, 71, 89) of 93. Roughly two epochs in total, on GB300 nodes.

Draft architecture, block size, sequence length, and loss weights are YAML-configurable — see the TorchSpec repo.

---

Quick Start

Requirements

For serving Kimi-K3 itself — hardware, parallelism and engine flags — follow the official vLLM recipe: recipes.vllm.ai/moonshotai/Kimi-K3.

Enable the draft

Add to your Kimi-K3 vllm serve command:

--speculative-config '{"method": "dspark", "model": "Inferact/Kimi-K3-DSpark", "num_speculative_tokens": 7, "attention_backend": "FLASHINFER_MLA", "draft_sample_method": "probabilistic", "rejection_sample_method": "block"}'