slippedJim/Kimi-K3-Draft

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.

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 — 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 — 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.

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.