speculative : capped recurrent planes for DFlash (--spec-draft-rs-planes)

[Piggidragon]

Lets DFlash use the capped recurrent planes that draft-mtp already has, renames the flag to match, and changes which states the capped planes keep.

What changes

• --spec-mtp-rs-planes becomes --spec-draft-rs-planes. The old name is still accepted, and the env var is now LLAMA_ARG_SPEC_DRAFT_RS_PLANES.
• draft-dflash can cap the target recurrent planes below spec-draft-n-max + 1. A rollback to a kept state selects its plane directly; any other one restores the pre-verification plane and replays the verification batch on the GPU, the same path MTP uses.
• DFlash has no replay state of its own: its draft state is the draft KV cache, which the replay already rewinds with seq_rm. The only DFlash-specific code is a replay-state hook that saves nothing.
• The capped planes no longer have to keep the newest states. --spec-draft-rs-planes-lead N sets how many planes keep the states after the first draft tokens. One plane always keeps the newest state, which a fully accepted draft continues from. The default picks the layout per mode, so MTP keeps its current behavior.
• ggml_gated_delta_net takes a leading_snapshots parameter for this. It is implemented in the CPU, CUDA and Vulkan kernels and in the CUDA cache fusion.
• The planes still belong to one speculative mode. Eagle3, DSpark, and combinations of two rollback modes are rejected as before.

Why the layout matters

The capped planes used to keep the states after the last tokens of the verification batch. That suits MTP, which accepts most of its draft. DFlash rejects early, so the states it needs were exactly the ones not kept.

Share of verifications by accepted draft count, uncapped, n-max 7, 1422 verifications:

accepted	0	1	2	3	4	5	6	7
share	24.5%	19.5%	15.0%	9.9%	7.7%	5.0%	3.6%	14.7%

Full acceptance is a second peak, which is why one plane stays on the newest state. With 4 planes the old layout covers accepted counts 5, 6 and 7 (23%), the new one covers 0, 1 and 7 (59%).

Results

Qwen3.8-27B-UD-Q4_K_XL + DFlash2-Q4_K_M over RTX 4070 + RTX 3060, host KV cache, parallel 2, n-max 7, 131k per slot, except where noted. Measured on a local merge of this PR with the open PRs this setup needs. "old" is lead 0, the layout before this PR; "new" is the default DFlash layout. Decode in t/s.

Correctness. Under tensor split and under layer split, single-request greedy generations with 4 and 5 planes, in both layouts, are byte-identical to the uncapped run and accept the same draft tokens. Two concurrent requests are not byte-identical to the uncapped run in any capped layout; see "Not covered".

Tensor split.

planes	kept states	replays	decode, short	decode @32k	peak 4070	peak 3060
8 (default)	all	0%	34.4	14.3	11587 MiB	11513 MiB
4, old	5, 6, 7	76%	18.4	8.3	11099 MiB	11023 MiB
3, new	0, 7	61%	20.5	8.9	10931 MiB	10857 MiB
4, new	0, 1, 7	40%	23.4	10.4	11129 MiB	11053 MiB
5, new	0, 1, 2, 7	26%	26.8	11.5	11249 MiB	11175 MiB

Layer split.

planes	replays	decode, short	decode @32k	peak 4070	peak 3060
8 (default)	0%	32.8	11.7	10417 MiB	11723 MiB
4, old	77%	20.1	6.9	9701 MiB	11229 MiB
3, new	61%	22.1	7.4	9543 MiB	11083 MiB
4, new	41%	24.9	8.5	9701 MiB	11229 MiB
5, new	28%	27.0	9.3	9943 MiB	11371 MiB

Single GPU, RTX 4070, Qwen3.8-27B-UD-IQ2_M + DFlash2-Q2_K, parallel 1, 65k context.

planes	replays	decode, short	decode @32k	peak
8 (default)	-	does not fit	-	-
4, old	81%	30.1	16.8	11615 MiB
3, new	61%	34.2	18.6	11489 MiB
4, new	41%	38.0	21.2	11615 MiB
5, new	27%	40.9	23.4	11771 MiB

Uncapped does not fit on one 12 GB card: 9228 MiB of weights plus 1197 MiB of planes leave too little, and the run dies allocating a 182 MiB compute workspace. There the cap is not a trade, it is what makes the configuration possible, and the new layout is 26% faster than the old one at the same peak VRAM. The two layouts produce identical output and identical acceptance there.

Each plane holds one recurrent state per sequence, about 150 MiB on this model, so the 4-plane runs hold 1197 MiB of recurrent state instead of 2394 MiB with two slots.

Prefill is unchanged in every run: 495-499 t/s tensor split, 536-544 t/s layer split, 903-904 t/s single GPU.

The decode cost tracks the replay share: a replay is a second target decode of the verification batch. With the new layout, 4 planes go from 54% to 68% of uncapped decode under tensor split, and 3 planes beat the old 4-plane layout while saving another 170 MiB per GPU.

Not covered

• DSpark and Eagle3 still cannot cap. DSpark shares the DFlash implementation and would likely work, but it is untested.
• Only CUDA was measured. The Vulkan kernel implements the new layout but was not run.
• MTP was not re-measured. Its default layout is unchanged, and its capped path is the one this PR started from.
• Two concurrent requests: no errors, greedy output is not byte-identical between plane counts. A replay step decodes only the replaying slot and so changes how the other slot is batched. Uncapped runs show the same thing: concurrent vs single request already diverges at a near-tie. The logprobs of accepted draft tokens are not reported, so the margin at those tokens was not measured.
• An adaptive layout, chosen from the acceptance history at runtime, is possible but not implemented. Changing the layout rebuilds the graph, so it would need hysteresis.

---

AI assistance: the implementation, tests and measurements were done with Claude Opus 5; each commit carries Assisted-by:.

🤖 Generated with Claude Code

https://claude.ai/code/session_01DPWnbaJURJQJbbQv6e4D8y