Kimi-K3-DSpark: MLA-native DSpark draft model for Kimi-K3

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,76 @@
+---
+license: mit
+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.**
+
+| Category | Benchmark | n | Mean Accept Length | Per-pos Accept Rate |
+|---|---|---:|---:|---|
+| Dialogue | MT-Bench | 80 | 2.57 | 0.624, 0.376, 0.222, 0.141, 0.096, 0.067, 0.048 |
+| Long-context | SPEED-Bench `throughput_16k` / low-entropy | 302 | 3.85 | 0.796, 0.609, 0.462, 0.353, 0.271, 0.208, 0.154 |
+| Code | SPEED-Bench · coding | 61 | 4.41 | 0.843, 0.666, 0.536, 0.440, 0.371, 0.309, 0.249 |
+| Multilingual | SPEED-Bench · multilingual | 47 | 2.85 | 0.650, 0.424, 0.280, 0.195, 0.139, 0.098, 0.066 |
+| Writing | SPEED-Bench · writing | 76 | 2.40 | 0.603, 0.342, 0.191, 0.115, 0.072, 0.046, 0.030 |
+| RAG | SPEED-Bench · rag | 55 | 2.30 | 0.591, 0.320, 0.167, 0.096, 0.061, 0.038, 0.025 |
+| QA | SPEED-Bench · qa | 80 | 2.06 | 0.549, 0.272, 0.124, 0.059, 0.030, 0.016, 0.009 |
+
+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"}'
+```

config.json

@@ -0,0 +1,54 @@
+{
+  "architectures": [
+    "K3DSparkModel"
+  ],
+  "model_type": "k3_dspark",
+  "hidden_size": 7168,
+  "intermediate_size": 14336,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 64,
+  "num_key_value_heads": 64,
+  "q_lora_rank": 1536,
+  "kv_lora_rank": 512,
+  "qk_nope_head_dim": 128,
+  "qk_rope_head_dim": 64,
+  "v_head_dim": 128,
+  "mla_use_nope": false,
+  "mla_use_output_gate": false,
+  "vocab_size": 163840,
+  "rms_norm_eps": 1e-05,
+  "max_position_embeddings": 1048576,
+  "rope_theta": 50000.0,
+  "num_target_layers": 5,
+  "target_hidden_size": 7168,
+  "target_num_hidden_layers": 93,
+  "target_layer_ids": [
+    2,
+    23,
+    47,
+    71,
+    89
+  ],
+  "mask_token_id": 163837,
+  "bos_token_id": 163584,
+  "eos_token_id": 163586,
+  "pad_token_id": 163839,
+  "markov_rank": 256,
+  "markov_head_type": "vanilla",
+  "enable_confidence_head": true,
+  "confidence_head_with_markov": true,
+  "tie_word_embeddings": false,
+  "draft_vocab_size": 163840,
+  "_torchspec_version": "0.1.0",
+  "torch_dtype": "bfloat16",
+  "rope_parameters": {
+    "rope_type": "yarn",
+    "factor": 32.0,
+    "original_max_position_embeddings": 32768,
+    "rope_theta": 50000.0,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26e41e0e322f366247a8fc300ef9d62497d236549491939aef3440ad5b0f76ae
+size 7124633450