perf(speculative): draft a tree instead of a chain, so an MTP round emits more than 2.33 tokens

[inureyes]

Why now

Phases 0 and 3 of #1185 have landed and the profile has been re-measured. The drafter side is no longer where the time is.

Round split on M5 Max, qwen3.8-27b-4bit with the qwen3_5_mtp drafter, --draft-block-size 3 --temp 0, after #1203:

component	ms/round	share
verify forward (T=3)	~40.0	~85%
draft_block	2.70	~5.7%
accept hook	2.74	~5.8%
finalize + residual	~1.6	~3.4%

A round costs about 47 ms and emits 2.33 tokens. Driving the drafter to zero cost leaves 41.6 ms and the same 2.33 tokens: 56.0 tok/s, or 1.77x classic decode. We are at roughly 1.5x. Everything left is within 0.27x, and none of it is on the drafter.

Past that ceiling the round has to emit more tokens, not cost less.

Deeper linear chains do not do it, and that is measured

The obvious move is a wider block. It fails, and not marginally. Sweeping --draft-block-size on this pairing (single sweep, thermally degraded host, so read the trend and not the absolute tok/s):

block	acceptance	emitted/verify	tok/s
3	0.65	2.3	21.30
4	0.49	2.47	16.24
6	0.35	2.72	12.12
8	0.25	2.74	9.12
12	0.17	2.82	5.80

Acceptance collapses by a factor of four while emitted-per-verify moves 2.3 to 2.8. A linear chain spends its whole budget on one guess about the future, so each extra position multiplies the chance that the prefix already diverged. The verify cost rises with the block and the yield does not.

That is the argument for branching rather than lengthening: spend the extra verify positions on alternatives at uncertain steps instead of on more of a single guess.

The headroom exists on the verify side

From the M5 Max report on #1182 (verify_forward_cost_scaling): this hardware verifies T=4 for 1.29x the price of T=1 and T=8 for 2.48x. A multi-token verify forward is close to free per extra position here, which is precisely the property tree drafting needs and the property a linear chain cannot convert into tokens.

Rough arithmetic with today's numbers: if a tree lifted emitted-per-round from 2.33 to about 3.2 while the verify forward grew to T=4 cost, the round lands near 48 ms for 3.2 tokens, about 66 tok/s, or 2.1x classic. That is a projection from the cost model, not a measurement, and the acceptance term is the part that has to be measured rather than assumed.

What this needs

Substantially more than Phases 0 through 4. Three pieces, none of which exist today:

1. A tree attention mask on the verify path. The verify forward currently builds a causal mask over a linear block. A tree needs each node to attend to its ancestors only. The mask is a function of the tree topology, which changes per round, so it cannot be cached the way the causal mask can.
2. A tree-aware accept walk. speculative_walk accepts the longest matching prefix of a chain. For a tree it has to descend the branch the target actually agreed with, which changes both the walk and what emitted_per_verify means.
3. Tree-aware rollback. verify_finalize trims the KV cache back to the accepted prefix. With a tree, the rejected siblings' positions were also written and have to be trimmed correctly. This is the piece most likely to produce silent corruption if it is wrong, and it interacts with the recurrent state in the GatedDeltaNet layers, which cannot be trimmed the way a KV cache can (see fix(speculative): Qwen 3.5 MTP output is not byte-identical to classic decode, and the exactness probe does not catch it #1201 for what that failure mode looks like in practice).

Where the branches go is a fourth decision: EAGLE-2 branches by drafter confidence, SpecInfer by a fixed topology. Confidence-based branching needs the drafter's logits, which the current step already computes and discards past the argmax.

Exactness

Unchanged in principle: the target verifies every node, so tree drafting can only cost acceptance, never output. But the byte-identity contract is a per-position claim against the single-token chain, and #1186 shows it depends on which MLX kernel the verify forward dispatches at each M. A tree changes M per round, so the #1189 probe has to be re-derived for the tree case rather than assumed to carry over.

Do not build this on the assumption that a passing probe implies end-to-end parity. It does not, which is what #1201 established the hard way.

Acceptance criteria

• Tree topology, mask construction, walk and rollback implemented behind a flag, defaulting off.
• Rollback correctness pinned by test, including the recurrent-state path, before any throughput number is quoted.
• Acceptance and emitted-per-round measured against the linear chain at the same block budget, on the Qwen 3.8 27B pairing and regression-checked on Gemma 4.
• Temperature-0 output compared against classic decode end to end, by diffing generated text, not by trusting the probe.
• The fix(speculative): gate MTP on a measured block-vs-chain exactness probe #1189 exactness gate extended to cover the tree case, or the tree explicitly gated to configurations where it holds.
• Measured on M5-class and on M1 Ultra, since the verify cost scaling that motivates this differs by generation.

Out of scope

The verify-side early-exit walk (#1179), which is a different way to spend the same verify forward and should be measured against this rather than combined with it.

Refs #1185, #1186, #1189, #1201

[inureyes]

Increment 1 landed on a branch: topology, mask, walk

feat/issue-1204-tree-drafting, commit d9650500. speculative::mtp::tree holds a DraftTree (nodes in topological order, parents[i] < i), its additive attention mask, and the accept walk. All pure, no device, no cache, no model, and nothing calls it yet.

Two equivalences carry the design and are what the tests pin:

• a linear tree masks exactly as a causal mask, for every block width from 0 to 5
• a linear tree walks exactly as speculative_walk, matching both the accept count and the emitted tokens at every accept length

Those are what let the tree path be switched on later without changing output on the configurations it subsumes.

One property worth stating because it bounds the rest of the work: the accepted set is always a path, never a subtree. A node can carry several children, but the target emits one argmax per position, so at most one child can match. Pinned by a test that gives two siblings the same token and checks the walk takes one and stops fanning out.

The rollback constraint, which decides feasibility

This came out of reading rollback_speculative_cache and it is not an implementation detail, so it should be on the record before anyone plans around the issue's original risk ranking.

Rollback today is two different mechanisms:

// attention: trim the tail
let trim = block_size - n;      // n = max(accepted) + 1
kv.trim(trim);

// GDN: replay the captured block over the accepted prefix
// (comment in-tree: "replay the captured block over `[:, :n]`")

Tail truncation is correct only because a linear accept set is a prefix. A tree's accepted path is scattered across the verified positions, so the cache has to keep those entries and compact them. KVCache exposes trim_to(new_len) and nothing else, so this needs a new primitive: either a gather along the sequence axis, or trimming the whole verified block and re-running a forward over the accepted path, which costs a small forward per round.

The recurrent side is in better shape than this issue assumed. The issue lists rollback as most likely to corrupt silently "and it interacts with the recurrent state, which cannot be trimmed the way a KV cache can". True, but Qwen 3.5 never trims it: it replays the captured block to rebuild the state. Replaying a gathered path instead of a prefix is an indexing change on the replay input, not a new mechanism. So the ranking inverts: the attention KV is the piece that needs new capability, and the GatedDeltaNet state is the piece that already generalizes.

What that means for the plan

The next increment should be the KV gather, not the drafter or the verify plumbing, because it is the only piece that can turn out to be infeasible. If a sequence-axis gather is cheap the rest is ordinary work; if it is not, the fallback is a re-run over the accepted path and the round budget has to absorb one extra small forward, which changes the arithmetic this issue is built on.

Nothing in this increment claims throughput, per the criterion that rollback correctness is pinned first.

[inureyes]

Redirection: the target this issue was measured on cannot verify a tree cheaply

Increment 2 (ecf2cdc4, fde88073) settled the rollback question: KVCache::gather_positions selects and compacts a scattered set of positions, and it works for Fp16, Int8, Turbo4, Turbo4Asym and Turbo3Asym. Only Turbo4Delegated refuses, for a layout reason (its V is split across packed cold storage and a hot ring at cold_offset), and pool-backed caches refuse for want of a dense buffer.

Worth recording that the first version of that gather inherited a Fp16 | Int8 restriction from trim_front_keep without checking whether its cause applies. It does not: that restriction exists because a head trim advances live_start instead of moving data and the Turbo fetch paths slice from 0, whereas the gather reindexes physically and leaves live_start at 0. Every dense mode keeps its per-token state on the same sequence axis trim already slices, so selecting a subset of that axis is the same operation with an index list instead of a bound.

The blocker is the recurrent layers, and it is correctness before it is cost

Qwen 3.8 27B, the pairing every number in this issue comes from, is a hybrid. full_attention_interval: 4 over 64 layers, and layer_types alternates three linear_attention to one full_attention: 48 GatedDeltaNet layers and 16 attention layers.

Attention verifies a tree the way this issue assumed. Flatten the nodes into one sequence, hand the tree mask to the forward, and siblings cannot see each other. That is 16 of the 64 layers.

The other 48 cannot. gated_delta is a recurrence, documented in-tree as a "sequential per-token loop" with a chunked scan that "computes the exact same recurrence". A fold has no mask. Flattening a tree through it carries state from a node into its sibling, which is wrong, and wrong silently.

So the choice is between:

• evaluating the recurrent layers once per leaf path, which multiplies 75% of the layers by the leaf count, in a design whose entire premise is that extra verify positions are nearly free
• a tree-aware scan kernel that forks state at branch points, which has to cover both the chain-parity kernel and the chunked scan
• restricting tree drafting to pure-attention targets

The measurement this issue leans on, T=4 for 1.29x the price of T=1, was taken on a flattened single sequence. It does not describe per-path recurrent evaluation, so it cannot be used to argue option one is affordable.

New order

Complete tree drafting on a pure-attention target first, measure there, and only then decide whether the GDN scan kernel is worth writing.

Gemma 4 is that target. gemma4.rs says so directly ("Gemma 4 has no SSM/GDN state"), and its caches are KVCache and RotatingKVCache with nothing recurrent. The pairing is available: gemma-4-12b-it-4bit with the gemma-4-12b-it-assistant-4bit drafter.

This reorders the issue rather than shrinking it. The kernel question stays open, but it becomes a question with an answer in front of it: whether trees pay at all on a target that can verify them at the price this issue assumed. If they do not, the kernel is moot; if they do, the measurement is the argument for writing it.

Re-measuring verify cost scaling on the Gemma 4 pairing now, since the 1.29x is a Qwen number and does not transfer.

[inureyes]

Linear widening on Gemma 4: it bends, and the bar for a tree is now a number

Measured on M5 Max, gemma-4-12b-it-4bit with the gemma-4-12b-it-assistant-4bit drafter, --temp 0 -n 120, cooled before the sweep and 20 s between arms. Classic decode on the same host is 43.3 tok/s.

block	acceptance	emitted/verify	proposals/round	verify ms/round	draft ms/round	round ms	tok/s	vs classic
2	0.803	1.803	1.00	23.80	1.11	24.92	72.95	1.68x
3	0.677	2.333	2.00	25.23	2.03	27.27	86.30	1.99x
4	0.585	2.705	2.95	26.20	2.95	29.16	93.54	2.16x
5	0.550	3.132	3.92	28.27	3.89	32.17	98.16	2.27x
6	0.426	3.051	4.87	32.10	4.80	36.90	83.38	1.93x
8	0.360	3.400	6.74	40.05	6.48	46.54	73.68	1.70x

Linear widening peaks at block 5 and declines after. So the curve does bend on this pairing, which is what the tree argument needs to be true; it just bends later and more gently than on Qwen, where widening from 3 to 12 took throughput from 21.30 down to 5.80 tok/s.

Extra verify positions are not free here either

23.80 ms at block 2 to 40.05 at block 8 is about 2.7 ms per extra position, roughly 11% of the base verify. The 1.29x-for-T=4 figure this issue quotes is a Qwen number and does not describe this pairing; on Gemma the same widening costs more, proportionally, than the headline suggested.

That matters because the tree's whole premise is that positions are cheap enough to spend on alternatives.

The bar a tree has to clear

At the linear optimum (block 5): the round is 32.17 ms for 3.132 emitted tokens, so 10.27 ms per token. One more verify position costs about 2.7 ms plus roughly 1.0 ms of drafting, so a branch position has to add at least about 0.36 emitted tokens per round to beat linear.

For comparison, the marginal linear position past the optimum is worth less than zero: going from 5 to 6 added a position and lost emitted-per-verify, 3.132 to 3.051, because acceptance fell faster than the extra depth paid.

So the experiment that decides this issue is now well posed: build the tree at a total node count matching block 5's four proposals, spend those nodes on branches instead of depth, and see whether emitted-per-verify clears 3.49. Below that, trees lose to simply setting --draft-block-size 5.

Two findings on the way, both separable from tree drafting

The adaptive controller costs about 5% on this pairing. effective_mtp_block_size pins the drafter at its configured depth (4 here) until recent accept lengths show that prefix is usually fully accepted, and at 0.585 acceptance it never expands. Block 5 is 98.16 tok/s against block 4's 93.54, so the policy is holding this pairing below its own optimum. The sweep above needed a temporary prefer_requested_block_size = true on the Gemma 4 drafter to get past it, which was reverted after measuring.

--draft-block-size is silently ignored when that happens, and the diagnostics hide it. Requesting 6 produced a run byte-identical to requesting 4, while the round-loop line still reported block_size=6. The diagnostics report the requested width, not the width actually used, so an operator tuning this flag has no way to see that the controller overrode them. Worth logging the effective width alongside the requested one.

[inureyes]

Increment 3: the mask is now the shape a forward actually takes

0d83f28d. additive_mask_with_history(offset, masked) produces [n, offset + n] following create_causal_mask's convention: query row q at logical position q + offset, the first offset columns the KV history every node attends in full, the last n following the topology.

The [n, n] form from the first increment was not usable against a real cache. It stays because it is what the topology tests compare against a plain causal mask.

Third equivalence pinned: a linear tree's history mask equals a causal mask built with the same offset, element for element. With the two already there (a linear tree masks as causal at offset 0, and walks as speculative_walk), that covers what the tree path has to reproduce before it can be turned on.

The injection seam is already open

Gemma4Model::forward_hidden_with_speculative_sinks takes mask: Option<&MlxArray> as its fourth parameter and threads it straight to forward_hidden_with_caches_and_speculative_sinks. The MTP adapter passes None today and builds nothing of its own.

So wiring the tree into verify is passing this mask instead of None, not adding a parameter or changing a signature on the model side. What remains is above it: MtpTarget::verify_forward has no way to carry a mask, and the round loop has no tree to build one from.

Remaining, in order

1. MtpTarget::verify_forward carries an optional mask; Gemma 4's adapter forwards it. Qwen 3.5's adapter rejects a tree outright rather than silently flattening one through 48 GatedDeltaNet layers.
2. The drafter produces a DraftTree instead of a Vec<i32>, behind a new trait method that defaults to wrapping the linear block, so nothing changes for drafters that do not implement it.
3. Round loop: build the tree, mask, verify, walk, and roll back through gather_positions, behind a flag defaulting off.
4. Measure against the bar: emitted-per-verify above 3.49 at the node count of block 5, or trees lose to --draft-block-size 5.

Step 1 is where a Qwen tree has to be refused explicitly. A flattened tree through a recurrence produces wrong state silently, and this issue should not ship a path that can do that by omission.