slippedJim/Kimi-K3-Draft

Kimi-K3 DSpark Draft (ATOM / ROCm)

An MLA-native DSpark draft model for Kimi-K3, served through ATOM's dspark speculative method on 8 × MI355X. 5 dense layers with non-causal attention drafting 7 tokens per pass, a low-rank sequential Markov head, and a confidence head. 68 tensors, 3,562,312,961 parameters, bf16.

Trained by off-policy prefill distillation: Kimi-K3's answers are generated once and replayed as prefill, so the teacher only publishes hidden states. The hidden states come from ATOM, the same engine that serves the draft.

---

Performance

tok/fwd = tokens per target forward step = 1 + accepted draft tokens / forward steps. Includes the bonus token the target emits every step, so 1.0 at zero acceptance and capped at 8. Same scale as the acceptance length reported by Inferact/Kimi-K3-DSpark.

ATOM, TP=8, fp8 KV, temperature=0, concurrency 1, max_tokens=12288, official prompt counts, Kimi-K3 chat template.

benchmark	prompts	tok/fwd	acceptance	reference card (vLLM)
GSM8K	1319	4.59	51.3%	5.64
MATH-500	500	3.53	36.2%	—
AIME 2026	30	2.57	22.5%	2.72
HumanEval	164	3.94	42.1%	5.34
MBPP	256	3.49	35.6%	—
MT-Bench	80	2.92	27.5%	3.14
SWE-bench Pro	128	2.95	27.8%	—
SPEED-Bench coding	80	4.50	50.0%	—
SPEED-Bench multilingual	80	3.57	36.8%	—
SPEED-Bench rag	80	3.42	34.6%	—
SPEED-Bench qa	80	3.11	30.2%	—
SPEED-Bench writing	80	2.70	24.2%	—
SPEED-Bench low-entropy (16k)	512	1.80	11.5%	—
mean		3.32

Accepted-length distribution, GSM8K (fraction of forward steps by tokens accepted): k=0 19.4%, k=7 28.1%.

Against the reference draft, same machine

Inferact/Kimi-K3-DSpark run on the same node, image, protocol and prompt counts:

benchmark	this draft	reference draft, same stack	ratio
GSM8K	4.59	4.90	93.7%
MATH-500	3.53	3.44	102.7%
AIME 2026	2.57	2.76	93.1%
HumanEval	3.94	4.07	96.9%
MBPP	3.49	3.75	93.1%
MT-Bench	2.92	3.10	94.3%
SWE-bench Pro	2.95	3.47	84.9%
SPEED-Bench coding	4.50	4.57	98.3%
SPEED-Bench multilingual	3.57	3.48	102.8%
SPEED-Bench rag	3.42	3.56	96.2%
SPEED-Bench qa	3.11	3.07	101.4%
SPEED-Bench writing	2.70	2.82	95.6%
SPEED-Bench low-entropy (16k)	1.80	4.33	41.7%
mean	3.32	3.64	91.1%

The reference draft's readings on this stack differ from its published card by −24% to +2% depending on the benchmark, so the ratio is taken against the same-stack measurement.

---

Training data

Every answer regenerated by Kimi-K3 itself. Prompts deduplicated, image-bearing rows removed, benchmark prompts excluded by 13-gram matching.

source	rows
slippedJim/ATOM-regen-nine-category-prompt	5,030,682
slippedJim/ATOM_regen_seeklight_kimi_mtp (text-only rows)	327,467
merged and shuffled	5,358,149
surviving the 8192-token window	5,011,091

Composition of the merged set:

category	share
multilingual (ja/de/it/es/fr)	74.8%
chat	7.7%
STEM	6.5%
kimi-mtp	6.1%
aya	2.7%
math	1.7%
code	0.43%

Sequence length of surviving rows: mean 2009, p50 1399, p90 4560, max 8191.

Training configuration

epochs	1, from scratch
steps	39,148
global batch / micro batch	128 / 1
max sequence length	8192 (overlong rows dropped, not truncated)
lr	5e-5 peak, cosine, 4% warmup, min_lr 3.93e-5
block size (speculative tokens)	7
aux hidden layers	2, 23, 47, 71, 89 of 93, post-norm
anchor_num	512
loss	CE 0.1 / L1 0.9 / confidence 1.0, position decay exp(-k/4)
KV / index cache	fp8 (matches the serving path)
topology	4 × ATOM TP=8 teachers + 1 × 8-rank FSDP2 draft node
throughput	5.7 s/step, ~65 h wall clock, 0 restarts
final training-side eval AL	2.767

---

Quick Start

ATOM is pinned by digest: atom 0.1.6rc1.dev275, torch 2.13.0+rocm7.14.0, HIP 7.14.60850.

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). Acceptance counters are at /debug/mtp_stats. max_num_seqs stays at 8: K3's KDA recurrent state is allocated per slot, and at 64 the state pool alone asks for 28.9 GB against a ~19 GB KV budget.

---

Limitations

• Nine of the reference card's thirteen benchmarks are not measured for this release: MATH-500, MBPP, SWE-bench Pro, and six SPEED-Bench splits.
• Long-context acceptance is unmeasured. The serving window here is 16384 tokens.
• nvidia/OpenCodeInstruct is not in the training mix; the reference draft's data includes it.
• No training checkpoint is published.