[WIP, Roadmap] Speculative decoding: MTP and DFlash (tracking) #828

Tracking issue for speculative decoding in FreeToken. Related: #79 (roadmap).

WIP: the first version has landed. The issue stays open for the open items in Support by model and PRs, and is updated as they move.

Scope

• Draft methods:
  • MTP, from the checkpoint's own MTP layers, or for Gemma 4 from the gemma4-assistant sidecar (which keeps no KV and reads the target's).
  • DFlash 2 drafts. DFlash v1 drafts are refused.
  • DSpark is not started.
• Targets: see Support by model. DFlash is not tied to one model family: a target only has to export the hidden states of the layers its draft reads.
• On the Qwen families this covers GatedDeltaNet hybrids, MoE offload and hybrid modes, and multimodal prompts.

Design

• Layers. The engine and the scheduler see one runtime interface, SpecRuntime (engine/spec_protocol.py), implemented by SpecRunner (spec/runner.py). The runner holds one Proposer (spec/proposer.py), the draft side of a method. Each method has a subpackage under spec/ (spec/mtp/, spec/dflash/).
• Drafts. Draft networks live with their model family (models/qwen3_5_moe/mtp.py, models/gemma4/mtp.py), and a family registers them in ModelSpec.spec_drafts as method -> class. --spec-algo mtp runs MTPProposer, or SharedKVMTPProposer for a family that declares shared_kv_mtp. spec/ imports no model family.
• Flow. Each scheduler iteration runs one batch: a prompt chunk if one is waiting, otherwise a verify batch for the running requests, for which it calls propose while building the batch. The engine runs the target forward and then calls on_prefill or post_verify in the same step, and the scheduler drains the batch before building the next. The table below says what each proposer does in its hooks.

• Target -> draft. The target forward hands the draft its hidden states through batch.features (FeatureSink in models/blocks.py): final_hidden for MTP drafts, and layer_output for the layers a DFlash draft reads.
• Per-request state. A per-request state (SlotStateSpec), such as GatedDeltaNet's, needs a verify handler (SlotVerify). The verify forward only reads the live state and stages scratch; commit rolls the state to the accepted count. The engine refuses a model with a state that has no handler.
• Prefix cache.
  • A draft whose KV slot i depends on token i+1 (the MTP drafts) sets radix_lookahead, so radix keys pair (t_i, t_{i+1}). The SWA radix cache keeps plain token keys.
  • Draft state has to survive a prefix hit. DFlash does not keep its sliding-window K/V for that. It checkpoints the window's context rows (DFlashWindow, one hidden-width row per position), which is 1.5-6x smaller than the draft's K/V for the released drafts. The checkpoint rides the GatedDeltaNet slots as a SlotStateSpec, so a prefix hit restores it with the GatedDeltaNet state, and the draft's K/V ring is rebuilt from it.
• Sliding-window targets. The SWA pool floor budgets decode_rows_per_req rows per step, and each request frees its out-of-window slots once an interval has slid out. A proposer opts in with target_may_be_windowed.
• Attention backends. The verify needs a backend that captures graphs of k+1 query rows per request: BackendInfo.uniform_extend_graph is device for fa, triton and qsa_sparse, and host_plan for fi. triton runs verify rows through its prefill kernel today (perf(attention): run up to 8 query rows per request in the triton decode kernel #900).
• Limits. --tp 1, and the normal loop (no overlap scheduling).
• Outputs. Speculative decoding must not change outputs. tests/spec/test_verify_contract.py compares greedy runs token by token on tiny models, and the real checkpoints are compared the same way.

Adding speculative decoding to a model

Questions to answer before writing code:

1. Draft. Which method does the release use: the checkpoint's own MTP layers, a sidecar that reads the target's KV (Gemma 4 assistant), DFlash 2, or DSpark? Where do its weights come from, and does an existing proposer fit?
2. Target forward. Does the model run a verify batch (k+1 rows per request) correctly, including per-token side inputs such as positions, mrope and PLE? Does it hand batch.features what the draft reads?
3. Per-request state. Does every SlotStateSpec of the family (GatedDeltaNet, KDA, ...) have a verify handler?
4. KV and prefix cache.
   • Does the draft keep KV of its own (declared in budget), and does that KV need radix_lookahead?
   • Is the target sliding-window, or DSV4 window-tier?
   • Must draft state be restored on a prefix hit?
5. Attention. Does the family's attention backend have a uniform_extend_graph? dsa, dsv4_sparse, dsv41_sparse and m3_sparse do not yet. Which backend runs the draft's own attention?
6. Tests. Add a tiny-config case to test_verify_contract.py: once with the oracle proposer for the target side, and once with the real draft class, asserting the draft is used. Compare greedy outputs on the real checkpoint.

Support by model

✅ supported, ❌ not yet, — none published. DFlash 2 drafts are incoai's releases. Prefix cache is the cache the engine runs for --cache-type radix: hybrid_radix for GatedDeltaNet / KDA hybrids, swa_radix for sliding-window and DeepSeek-V4 window-tier models, plain radix otherwise.

A draft with a sliding window (every DFlash 2 release uses 2048 tokens) on a hybrid_radix target follows the Qwen families. The draft's window is not kept as sliding-window KV under a window radix: cached prefixes would each hold draft K/V, which costs too much VRAM on edge hardware. Instead, the rows its K/V is computed from are stored as a slot state in hybrid_radix, restored with the target's linear-attention state on a prefix hit, and the K/V is rebuilt from them. For a radix or swa_radix target, there is no design yet.

Not in the first version

• Overlap scheduling together with speculative decoding.
• Tree drafts (not planned).

PRs

• Engine and scheduler seams for multi-row decode steps (no behavior change): refactor: add engine and scheduler seams for multi-row decode steps #833
• radix_lookahead keys, so that a prefix hit never reuses an MTP draft slot written for a different next token: feat(scheduler): add radix_lookahead keys to the prefix cache #834
• Linear-attention metadata in a LinearAttnBackend (no behavior change): refactor(attention): move linear-attention metadata into a LinearAttnBackend #837
• One shared GatedDeltaNet layer for the qwen3_5_moe and qwen4_exp families (no behavior change): refactor(layers): share one GatedDeltaNet op in layers/linear_attn #841
• Target-side verify: GatedDeltaNet verify and replay, rejection sampling, per-row presence and frequency penalties: feat(spec): verify draft tokens on the target side #836
• MTP for the Qwen3.5 family (HF checkpoints) and Qwen3.8-Flash-Next (HF checkpoint, or GGUF with a Strata MTP pack): feat(spec): add MTP drafts for the Qwen3.5 family and Qwen3.8-Flash-Next #844
• DFlash 2 drafts, with exact draft-context restore on prefix-cache hits: feat(spec): add DFlash 2 speculative decoding #857
• GGUF MTP weight loading for the Qwen3.5 family (side branch): feat(gguf): load mtp drafts for the qwen3.5 family from nextn blocks and mtp-*.gguf sidecars #856
• --speculative-token-map for a smaller MTP draft head: feat(spec): add --speculative-token-map for a smaller mtp draft head #892
• Gemma 4 MTP from the gemma4-assistant GGUF sidecar (dense, E2B / E4B, 26B-A4B), and with it sliding-window targets: feat(gemma4): speculative decoding with the assistant draft #895
• [Perf] Up to 8 query rows per request in the triton decode kernel, so verify steps stop falling back to the prefill kernel at long context: perf(attention): run up to 8 query rows per request in the triton decode kernel #900
• [Perf] CUDA graphs for the Gemma 4 assistant draft, which runs eager
• DFlash with prefix reuse on radix and swa_radix targets (no design yet)

PRs will be linked here as they open.