README.md

---
license: other
library_name: transformers
tags:
  - speculative-decoding
  - dspark
  - kimi-k3
  - draft-model
base_model: moonshotai/Kimi-K3
---

# Kimi-K3 DSpark Draft (ATOM / ROCm)

An **MLA-native DSpark** draft model that accelerates Kimi-K3 on **ATOM**, AMD's
inference engine, served through ATOM's `dspark` speculative method on 8 × MI355X.

The draft is trained by **off-policy prefill distillation**: Kimi-K3's answers are
generated once, cached, and then replayed as prefill so the teacher only has to
publish hidden states — no decode sweep in the training loop. The hidden states
come from ATOM itself, the same engine that serves the draft, so the numerics it
learns from are the numerics it meets at inference.

**DSpark** = a block-diffusion backbone of 5 dense layers with non-causal
attention, drafting 7 tokens in one parallel pass, a low-rank sequential Markov
head for intra-block dependency, and a confidence head. Mirroring Kimi-K3's MLA
attention means draft and target share one KV layout, so the draft's pages unify
with the target's cache. 68 tensors, 3,562,312,961 parameters.

---

## Performance

Acceptance length with 7 speculative tokens. `tok/fwd` = tokens produced per
target forward step = `1 + accepted draft tokens / forward steps`, including the
bonus token the target emits every step, so it is 1.0 at zero acceptance and
caps at 8. Same scale as the "acceptance length" reported by
[Inferact/Kimi-K3-DSpark](https://huggingface.co/Inferact/Kimi-K3-DSpark).

Measured on ATOM, TP=8, fp8 KV, `temperature=0`, concurrency 1, official prompt
counts, Kimi-K3 chat template enabled.

| benchmark | temperature=0 | acceptance | prompts |
|---|---|---|---|
| GSM8K | **3.96** | 42.3% | 1319 |
| HumanEval | **3.45** | 35.0% | 164 |
| MT-Bench | **2.81** | 25.9% | 80 |
| AIME 2026 | **2.45** | 20.7% | 30 |
| **mean** | **3.17** | | |

### Against the reference draft, same machine

Published acceptance numbers are engine-dependent, so rather than compare against
a card measured on other hardware, `Inferact/Kimi-K3-DSpark` was downloaded and
run on **the same node, image, protocol and prompt counts**:

| benchmark | this draft | reference draft, same stack | ratio |
|---|---|---|---|
| GSM8K | 3.96 | 4.90 | 80.8% |
| HumanEval | 3.45 | 4.07 | 84.8% |
| MT-Bench | 2.81 | 3.10 | 90.7% |
| AIME 2026 | 2.45 | 2.76 | 88.6% |
| **mean** | **3.17** | **3.71** | **85.4%** |

The reference draft's own readings on this stack differ from its published card
by −24% to +2% depending on the benchmark, which is why the ratio above is taken
against a same-stack measurement rather than against published numbers.

---

## Training

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

* [slippedJim/ATOM-regen-nine-category-prompt](https://huggingface.co/datasets/slippedJim/ATOM-regen-nine-category-prompt)
  — a nine-category pool from the NVIDIA **Nemotron** post-training collections
  plus CohereLabs/aya_dataset, spanning chat, code, math, STEM and five
  languages. Prompts are deduplicated, image-bearing samples are removed, and
  evaluation-set prompts are excluded by 13-gram matching.
* [slippedJim/ATOM_regen_seeklight_kimi_mtp](https://huggingface.co/datasets/slippedJim/ATOM_regen_seeklight_kimi_mtp)
  — general instruction prompts, used for the preceding epoch.

**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. Two epochs in total, on 8 × MI355X nodes:
four resident ATOM TP=8 teachers streaming hidden states to one 8-rank FSDP2
draft node.

Global batch 128, peak lr 5e-5 cosine with 4% warmup, KV and index cache in fp8
to match the serving path.

---

## Quick Start

ATOM is pinned by digest because `rocm/atom-dev:latest` is a rolling tag and has
moved in ways that deadlock K3's hidden-state extraction. The digest below is
atom `0.1.6rc1.dev275`, torch `2.13.0+rocm7.14.0`, HIP `7.14.60850`.

```bash
docker run -d --name atom-dspark \
  --device=/dev/kfd --device=/dev/dri --group-add video \
  --security-opt seccomp=unconfined --cap-add=SYS_PTRACE \
  --ipc=host --shm-size 128g --network host \
  -v /path/to/Kimi-K3:/target:ro -v /path/to/this/repo:/draft:ro \
  rocm/atom-dev@sha256:2f8bd4206ad15d014ae48115eae1ee9f1db83781848a8542de7177cfbd4ac914 \
  python -m atom.entrypoints.openai_server \
    --model /target --served-model-name Kimi-K3 \
    --method dspark --draft-model /draft --num-speculative-tokens 7 \
    --kv_cache_dtype fp8 -tp 8 --trust-remote-code \
    --max-model-len 16384 --max-num-seqs 8 --max-num-batched-tokens 10240 \
    --gpu-memory-utilization 0.93 --block-size 128 \
    --no-enable_prefix_caching --server-port 8000
```

The server log should show `Detected MLA DSpark drafter` and
`DSparkProposer aux capture on target layers: (2, 23, 47, 71, 89)`. Live
acceptance counters are at `/debug/mtp_stats`.

---

## Limitations

* **GSM8K is saturated for this architecture** — it did not move between the last
  two epochs despite 74k math prompts being added. Arithmetic derivations are
  template-like enough that 5 layers drafting 7 tokens appears to be the binding
  constraint, not data.
* **Code is the thinnest slice of the training mix.** The Nemotron code split
  collapses from 175k rows to 31k unique prompts after deduplication, and
  `nvidia/OpenCodeInstruct` was not included. HumanEval is correspondingly the
  weakest relative result.
* **Long-context acceptance is unmeasured.** The serving window here is 16384
  tokens, so benchmarks with 70k+ token prompts were not run.
* `training_checkpoint/checkpoint_3358.pt` belongs to the **preceding** epoch, not
  to the weights in this repo. Resuming training from it continues the earlier
  model.