slippedJim/K3_draft

Kimi-K3-DSpark-SDDD

A DSpark draft model for speculative decoding with moonshotai/Kimi-K3, trained with on-policy speculative distillation.

The file layout, tensor names and config.json are identical to Inferact/Kimi-K3-DSpark, so this is a drop-in replacement for it — only the weights differ.

Architecture

3,562,312,961 parameters in bfloat16 across 68 tensors:

Component	Shape	Note
embed_tokens.weight	163840 x 7168	frozen, copied from Kimi-K3
context_proj.weight	7168 x 35840	fuses the 5 captured teacher layers
5 x decoder layer	—	MLA, q_lora_rank=1536, kv_lora_rank=512
markov_head.markov_w{1,2}	163840 x 256	low-rank output head
confidence_head.proj	1 x 7424	per-position confidence

The draft consumes auxiliary hidden states from teacher layers [2, 23, 47, 71, 89] (93 layers total) and predicts block_size=7 positions ahead. There is no lm_head; embed_tokens is frozen and shared with the teacher, so only 2,387,907,841 parameters were actually trained.

Training

On-policy speculative distillation: the teacher generates its own continuations, and the draft is trained to predict the teacher's next tokens from the teacher's hidden states. This keeps the training distribution equal to the deployment distribution, which offline distillation on a fixed corpus does not.

Framework	LumenRL, batch-alternating scheduler
Teacher backend	ATOM AsyncLLMEngine, tensor parallel 8
Hardware	1 node, 8 x AMD Instinct MI350
Steps	5413
Global batch / micro batch	64 / 1
Learning rate	7.5e-5 cosine to 7.5e-6, 4% warmup
Gradient clipping	0.5
Weight decay	0
Precision	bfloat16
Rollouts	greedy, up to 992 new tokens, reasoning traces enabled
Data	476,904 conversations

Teacher rollouts and draft updates alternate in rounds of 50 batches: the teacher engine generates and captures hidden states, then releases the GPUs so the draft can train on the cached states.

Results

Simulated acceptance length over 7 speculative positions, measured on a held-out set of 32 prompts:

Step	Acceptance length	Eval loss
1700	1.1491	—
5100	1.86586	2.41002
5200	1.87052	2.39803
5300	1.88224	2.39511
5400	1.88704	2.37878

Gains had flattened to roughly +0.005 per 100 steps by the end of the run, in step with the cosine schedule reaching its floor.

Note on the metric: 32 prompts is a small evaluation set. Early in training the round-to-round spread was around 0.10, comparable to the signal; the late-training numbers above are far more stable but should still be read as approximate.

Usage

Load it exactly as you would Inferact/Kimi-K3-DSpark — same model_type (k3_dspark), same tensor names, same config keys. The draft requires a runtime that can expose the teacher's auxiliary hidden states at layers [2, 23, 47, 71, 89].

License

Inherits the license of the Kimi-K3 base model.