lightseekorg/kimi-k3-dflash2

Kimi-K3 DFlash2 draft (5-layer, lookahead-8, MLA + SWA)

Draft model for speculative decoding on moonshotai/Kimi-K3, served with TokenSpeed (--speculative-algorithm DFLASH).

Architecture

UPDATE: This repository now hosts a five-layer draft.

• DFlash2DraftModel, 5 layers, hidden size 7168, bfloat16
• block_size: 8 (lookahead-8), attention_mode: mla
• target_layer_ids: [19, 37, 66, 78, 90] - five auxiliary hidden-state taps
• layer_types: 4x sliding_attention (window 4096) + 1x full_attention
• yarn RoPE, rope_theta: 50000.0, factor: 32.0
• 7.39 GB of weights

Results

Full tables and caveats in benchmarks/RESULTS.md; raw per-row records in benchmarks/.

• Ten-benchmark matrix at production sampling (temperature=1.0, top_p=0.95, reasoning effort max): mean acceptance length 3.886, ranging from 5.851 on GSM8K to 2.657 on AIME 2026.
• Long context: acceptance plateaus near 4.8 from roughly 16k tokens onward and holds at 4.854 with 130,151 input tokens, with needle retrieval 1.000 at every length tested.

Please read the caveats before quoting these figures - in particular, the long-context numbers measure full-context MLA compute rather than sliding-window-equivalent inference.

Serving

python3 -m tokenspeed.cli serve <Kimi-K3> \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path <this repo> \
  --speculative-num-steps 7 --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 8 \
  --speculative-draft-model-quantization unquant \
  --attention-backend mla --drafter-attention-backend mla \
  --sampling-backend flashinfer --trust-remote-code \
  --tensor-parallel-size 8 --moe-backend flashinfer_trtllm \
  --max-model-len 32768 --gpu-memory-utilization 0.90 \
  --max-num-seqs 4 --max-cudagraph-capture-size 4 \
  --cudagraph-capture-sizes 1 2 3 4 \
  --disable-prefill-graph --disable-kvstore

--sampling-backend flashinfer is required for sampled decoding: it selects rejection-based verification, which consumes the draft's own distribution along with the request's temperature and top_p.

At startup the engine reports the auxiliary taps it captured; confirm they match (19, 37, 66, 78, 90).