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

Tracking issue for speculative decoding in FreeToken. Related: #79 (roadmap), #421 (MTP for qwen4_exp).

WIP: the scope, design and plan below may change as the PRs are developed and reviewed; this issue will be updated to match.

Scope

• Draft methods: MTP (the checkpoint's own MTP layers) first, then DFlash (DFlash and DFlash2 draft checkpoints).
• MTP targets: the Qwen3.5 / 3.6 / 3.8 family, dense and MoE, and Qwen3.8-Flash-Next (qwen4_exp). This includes GatedDeltaNet hybrids, MoE offload and hybrid modes, and multimodal prompts.
• DFlash is not tied to one model family: a target only has to export the hidden states of the layers its draft reads.

Design

• The engine and the scheduler see one runtime interface (engine/spec_protocol.py). Method code lives under python/freetoken/spec/, one subpackage per method (spec/mtp/, spec/dflash/).
• Draft networks live with their model family (for example models/qwen3_5_moe/mtp.py). They add no import between model families, and spec/ imports no model family.
• A verify step is a uniform extend of k+1 rows per request. GatedDeltaNet state is only read during the verify; the commit replays the accepted rows.
• Speculative decoding must not change outputs. Greedy runs with and without it are compared token by token on bf16 checkpoints.

Not in the first version

• Targets with sliding-window KV (gpt-oss, Gemma 4) and DeepSeek-V4 window tiers.
• Targets with KDA linear attention (GLM-5.3-Flash, glm5_next).
• Overlap scheduling together with speculative decoding.
• Tree drafts (not planned).

Plan

Expected merge order. An arrow points from a PR to one that needs it merged first.

Seams ──────────────────────────────┐
LinearAttnBackend ──────────────────┼──► Verify ──► MTP: Qwen3.5 + Flash-Next ──► DFlash
Shared GatedDeltaNet layer ─────────┘                 ▲     │
radix_lookahead keys ─────────────────────────────────┘     └──► GGUF MTP weight loading: Qwen3.5 (side branch)

The two refactors moved the linear-attention metadata into a LinearAttnBackend and merged the two GatedDeltaNet ops into one layer, so the verify path is written once for both model families. MTP for the Qwen3.5 family and for Qwen3.8-Flash-Next is one PR. GGUF MTP for the Qwen3.5 family is a side branch: it only adds weight loading (the nextn blocks in a target GGUF, and separate mtp-*.gguf files), and can merge any time after the MTP PR. DFlash comes after the MTP PR because that PR adds the rest of the qwen4_exp verify path (QSA verify graphs and the PLE verify handlers).

• 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 and DFlash2 drafts, with exact draft-context restore on prefix-cache hits
• GGUF MTP weight loading for the Qwen3.5 family (side branch)

PRs will be linked here as they open.