| --- |
| license: other |
| license_name: kimi-k3 |
| library_name: vllm |
| base_model: moonshotai/Kimi-K3 |
| pipeline_tag: text-generation |
| tags: |
| - dspark |
| - speculative-decoding |
| - draft-model |
| - mla |
| - vllm |
| - torchspec |
| - kimi-k3 |
| --- |
| |
| ## Model Overview |
|
|
| **Inferact/Kimi-K3-DSpark** is an **MLA-native DSpark** draft model that accelerates inference of [Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) on **[vLLM](https://github.com/vllm-project/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](https://github.com/lightseekorg/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](https://huggingface.co/datasets/philschmid/mt-bench) (80 dialogue prompts) and [NVIDIA SPEED-Bench](https://huggingface.co/datasets/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`](https://huggingface.co/datasets/lightseekorg/kimi-mtp-dataset) — general instruction prompts |
| - [`nvidia/OpenCodeInstruct`](https://huggingface.co/datasets/nvidia/OpenCodeInstruct) — single-turn coding |
| - A nine-category prompt pool from the NVIDIA **Nemotron** SFT/RL collections and [`CohereLabs/aya_dataset`](https://huggingface.co/datasets/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`, `max_seq_length=16384`, 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](https://github.com/lightseekorg/TorchSpec). |
|
|
| --- |
|
|
| ## Quick Start |
|
|
| ### Requirements |
|
|
| For serving Kimi-K3 itself — hardware, parallelism and engine flags — follow the official vLLM recipe: **[recipes.vllm.ai/moonshotai/Kimi-K3](https://recipes.vllm.ai/moonshotai/Kimi-K3)**. |
|
|
| ### Enable the draft |
|
|
| Add to your Kimi-K3 `vllm serve` command: |
|
|
| ```bash |
| --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"}' |
| ``` |
|
|
