DFlash2 speculative decoding for Qwen3.8-27B

[spokvulcan]

Proposed changes

Adds DFlash2 speculative decoding for the Qwen3.5 family. Closes #606.

Why this model in particular, Qwen3.8-27B scores 52 on the Artificial Analysis Intelligence Index (https://artificialanalysis.ai/models/qwen3-8-27b), the same as GPT-5.6 Luna and 14 points behind Claude Fable 5.1 at 66.

DFlash2 is the block drafter from inco.ai and z-lab (https://inco.ai/blog/dflash2/, reference code at https://github.com/z-lab/dflash). It's a separate small model, not a head inside the checkpoint: for Qwen3.8-27B it has five layers, reads the target's hidden states after layers 5, 19, 33, 47 and 61, and proposes 7 tokens at once, which the target then checks in a single forward pass. The drafter checkpoint is https://huggingface.co/incoai/Qwen3.8-27B-DFlash2. It has no embedding or output head of its own and borrows the target's. Output is token-for-token what the target would produce by itself.

What's in the PR

A drafter model and a speculative token iterator that follow the existing MTP drafter pattern (drafter protocol, factory, registry and generate overloads), plus a verify pass on the Qwen3.5 text model that checks a whole 8-token block in one forward pass and commits the caches to whatever prefix was accepted.

Numbers

Measured on a MacBook Pro 16-inch, Apple M3 Max (16-core CPU, 40-core GPU), 48 GB, macOS 26.5.2, with Qwen3.8-27B quantized to 4 bits (group size 64) and the drafter above, also 4-bit. Greedy decoding, block size 8, 192 generated tokens per run, thinking enabled. One prompt per category rather than the full benchmark sets, so treat the right column as a spot check.

"Acceptance length" is the average number of tokens each verify pass yields, counting the one token the target produces itself: 1.0 is plain decoding, 8.0 is every proposal accepted. The comparison column is Table 4 of the blog post linked above, same model and drafter, block size 8.

Task	inco.ai blog post	This PR	Prompt used
GSM8K	5.46	5.33	two grade-school word problems
MATH-500	5.28	5.82	three algebra and geometry problems
HumanEval	4.39	3.76	write a longest-common-subsequence function with tests
MBPP	4.79	3.76	same code prompt
MT-Bench	4.10	2.59	write a travel blog post about Hawaii
long prose	not reported	2.53	summarize a 6,000-token design document

Decode speed on that machine. The blog post doesn't publish tokens per second, so no comparison column:

Prompt	Plain TokenIterator	DFlash2	Speedup
math word problems	22.9 tok/s	77.0 tok/s	3.4x
algebra and geometry	22.9 tok/s	84.3 tok/s	3.7x
code	22.7 tok/s	54.6 tok/s	2.4x
travel blog post	21.7 tok/s	28.5 tok/s	1.3x
6,000-token summary	21.4 tok/s	32.8 tok/s	1.5x

Math and code match the blog post, and the z-lab reference implementation (Python/MLX) gives the same acceptance on the same prompts. Chat and prose sit below the MT-Bench number: with thinking on, most of the 192 tokens are the model's planning text, which the DFlash paper measured at roughly a quarter lower acceptance than the answer itself, and a 4-bit target loses a few more points against the unquantized one the published numbers use. Every row is token-for-token identical to plain TokenIterator output, and greedy acceptance is stable run to run.

Tests

Tests/MLXLMTests/DFlash2Tests.swift, all with tiny synthetic models so they run in the normal unit suite: config decoding and rejection of unsupported shapes, the context cache (placeholder rows, compaction), the visibility mask, the dynamic convolution against a naive loop, the selector, compiled versus eager drafter output, gated-delta replay against a prefix recompute, rejection sampling, sampling filter order against TopPSampler, and the iterator against mock target and drafter models: accept and bonus bookkeeping, penalties matching TokenIterator token for token, warm start over a prefilled cache matching a cold run, finalize rewinding undrained drafts, and the errors for unsupported inputs.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

AI usage

• I have read this PR description in full and approve it as my own, and it
  accurately describes the code changes.
• AI usage disclosure: written with Claude Code from the reference
  implementation and our earlier port, reviewed and benchmarked by me line
  by line.

[spokvulcan]

One more commit. stackSameInputProjections(in:) iterated modules(), whose array holds every projection module, so each block's originals stayed alive until the loop ended and the transient over a load was the sum of every stacked block instead of one. The loop now keeps only the stacking modules. testSameInputStackingReleasesEachBlockBeforeTheNext stacks eight quantized MLP blocks and bounds the peak at two; it held all eight before.

Qwen3.5-27B 4-bit in the app that embeds this library, MLX peak memory in GB:

phase	before	after
peak during projection stacking	22.97	17.84 (flat over the 15.9 GB model)