perf(speculative): cut the MTP drafter's per-step cost, now the binding constraint on M5-class hardware

[inureyes]

Summary

On M5-class hardware the MTP speculative round is no longer bound by the target's multi-token verify forward. It is bound by the drafter. A single drafter step costs about 10.5 ms against a full 27B target decode step of 30.3 ms, while the drafter only touches 1.56 GB of weights against the target's 16.05 GB. The drafter therefore runs at roughly 149 GB/s effective memory bandwidth where the target achieves 530 GB/s, a 3.6x efficiency gap that is pure overhead.

Closing that gap is worth an estimated 1.19x to 1.8x on the measured Qwen 3.8 27B pairing without touching acceptance or exactness. Going past that needs the algorithmic work in the last section.

Measured on Apple M5 Max 128 GB, macOS 26.6.1, branch feature/issue-1165-qwen3-5-mtp-drafter at ef3aa95e (#1182), models/qwen3.8-27b-4bit + models/qwen3.8-27b-mtp-bf16, offline CLI, 300 tokens, temperature 0, Time Machine stopped, 90 s cooldown between runs.

This applies to Gemma 4 MTP as well: the round loop, the Drafter trait hooks, and the sampling helper are shared. The numbers below are from the Qwen pairing because that is where the profile was captured.

The measured split

Classic decode: 32.96 tok/s, 30.3 ms/token.
MTP block 3: 39.37 tok/s, 1.19x. MTP block 4: 33.87 tok/s, 1.03x.

Per-round split from the round-loop diagnostics (block 3, 133 rounds, 2.248 tokens per round, 57.35 ms per round):

phase	ms/round	share
verify forward (T=3)	35.27	61.5%
draft_block (1 drafter step at this block size)	10.70	18.7%
accept hook + sampling + argmax (residual, uninstrumented)	11.22	19.6%
verify finalize (rollback)	0.16	0.3%
speculative walk, shared-kv re-arm	<0.01	~0%

Block 4 (126 rounds, 2.373 tokens per round, 71.88 ms per round): verify 39.95 ms (55.6%), draft_block 20.88 ms over 2 steps (29.0%), residual 10.85 ms (15.1%), finalize 0.20 ms.

The per-step drafter cost is flat at about 10.5 ms: block 3 runs one step inside draft_block for 10.70 ms, block 4 runs two for 20.88 ms (10.44 each). Per round there are two drafter forwards in total, one in draft_block and one in accept_verified_tokens.

The comparison that identifies the problem

Same pairing, same protocol, block 4, against the M1 Ultra record in docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md:

phase	M1 Ultra	M5 Max	change
verify forward (T=4)	137.6 ms	39.95 ms	3.44x faster
draft (per round)	19.9 ms	20.88 ms	no change
accept hook	~13 ms	10.85 ms	1.2x faster
round total	173 ms	71.88 ms	2.41x faster

The drafter step did not get faster at all across a GPU generation that made the target's multi-token forward 3.44x cheaper. Whatever dominates the drafter step is not memory bandwidth and not arithmetic.

The cost model

Per drafter step the weights actually read are:

tensor set	bytes
drafter weights (15 tensors, f16)	849 MB
borrowed target lm_head (weight 635.7 + scales 39.7 + biases 39.7)	715.2 MB
embedding row lookup	negligible
total	1.564 GB

At the target's demonstrated 530 GB/s that is 2.95 ms. Measured 10.5 ms. About 7.5 ms per step is overhead, and there are two steps per round, so roughly 15 ms of a 57.35 ms round, or 26%, is drafter overhead.

Note also that the borrowed LM head is 46% of the drafter's memory traffic. Qwen 3.8's vocabulary is 248,320 and the target is untied, so every drafter step pays a 248,320 x 5,120 projection to produce a single token id.

What the code does per drafter step

From src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs:

1. forward_hidden_stack uploads the token ids host to device via ffi::from_slice_i32 (line 356).
2. Embedding lookup, two pre-fc norms, a concat, the fc projection, one full decoder layer, final norm.
3. project_logits runs the full 248,320-wide LM head (line 389).
4. sample_one calls ffi::fused_sample, then ffi::eval followed by ffi::item_i32 (lines 400-411).

Step 4 is a full GPU pipeline flush plus a device-to-host scalar readback, and step 1 is a host-to-device upload, so there is a complete host round trip inside the autoregressive draft loop, once per drafted token. set_seed_from_hidden does the same thing again in the accept hook. At block 3 that is two flushes per round; at block 4, three.

This is the same class of problem #1179 identifies on the verify side ("projects one position at a time across the cxx bridge"), on the draft side instead.

Proposed work

Phase 0: attribute the 7.5 ms before optimizing

Nothing below should be built on the cost model alone. Add timers inside the drafter step and report ms for: id upload, layer forward, LM head projection, fused_sample, eval + item_i32 readback. Also add a accept_hook_ms counter to MtpDiagnostics, which today has no counter at all, so the accept hook is invisible and lands in the residual.

Acceptance: a diagnostics line that accounts for the full 10.5 ms and says which of the five components dominates.

Phase 1: remove the host round trip from the draft loop

Keep the drafted token id on device across the block. fused_sample already produces a device array; feed it straight into the embedding lookup for the next step instead of round-tripping through item_i32 and from_slice_i32. Materialize the whole block of ids in one eval at the end of draft_block, which is the pattern argmax_from_hidden_positions already uses on the verify side per #1179.

This requires an embedding forward that accepts a device id array, and a forward_hidden_stack variant that takes ids as &MlxArray rather than &[i32].

Expected: removes K-1 pipeline flushes per round at block K. Cannot change output, because the drafted ids are identical either way and every draft is verified by the target regardless.

Phase 2: stop paying the full-vocab LM head per drafter step

Two independent options, both safe for exactness:

2a. Fuse projection and argmax. At temperature 0 the 248,320-wide logits vector is materialized only to take its argmax. A fused LM-head-plus-argmax kernel never writes the logits. Saves a write and a read of the logits buffer plus a dispatch, and lets the projection stay in registers.

2b. Reduced draft vocabulary. This is what EAGLE-3 does with draft_vocab_size plus a d2t index mapping, and it is the single largest term available: restricting the drafter's projection to a static high-frequency subset of, say, 32,768 ids cuts 715 MB to about 92 MB per step. Derive the subset offline from token frequency over a calibration corpus and ship it beside the drafter, falling back to the full head when absent.

The correctness argument for 2b is important and worth stating explicitly: a draft that is wrong is rejected by the target, so drafter numerics cannot affect output. The only exposure is acceptance rate, which is measurable. If the true next token falls outside the subset, that round simply accepts fewer drafts.

Phase 3: quantize the drafter

The drafter ships bf16 and loads as f16, 849 MB. A 4-bit conversion brings it to roughly 250 MB including scales and biases. Combined with 2b a drafter step reads about 342 MB instead of 1.564 GB.

Same correctness argument as above: the drafter proposes, the target verifies, so quantization risks acceptance and nothing else. Needs a converted checkpoint (mlxcel-surgery is the natural home) and an acceptance A/B against the f16 drafter before it becomes a default.

Phase 4: cheap structural cleanups found while reading

• accept_verified_tokens builds hidden_cat with a chain of concatenate calls (model.rs:659-674). The positions pushed are always the contiguous range [keep, accepted] of verify_hidden, so this can be one slice instead of n concatenations and n-1 intermediate buffers.
• forward_hidden_stack rebuilds a causal mask on every multi-token call. It is a function of (s, cache_offset) and can be cached.

What this is projected to buy

Holding acceptance, block size, and the verify forward at their measured values and varying only the drafter step cost:

scenario	drafter step	round	tok/s	vs classic
measured today	10.5 ms	57.35 ms	39.4	1.19x
Phase 1+4, drafter at target bandwidth efficiency	2.95 ms	~42 ms	53	~1.61x
plus Phase 2b (32k draft vocab)	1.78 ms	~40 ms	57	~1.71x
plus Phase 3 (4-bit drafter)	0.65 ms	~38 ms	60	~1.82x
verify-forward-only ceiling	0	35.27 ms	63.7	1.93x

These are projections from the cost model, not measurements. They are included to rank the phases, and Phase 0 exists to test them.

Going past 1.93x

The table above converges on a hard ceiling: at 2.248 emitted tokens per round and a 35.27 ms verify forward, no amount of drafter optimization gets past 1.93x. Beyond that the round has to emit more tokens, not cost less.

The relevant measurement is that this hardware verifies T=4 for 1.29x the price of T=1, and T=8 for 2.48x (from verify_forward_cost_scaling, see the M5 Max report on #1182). So there is real headroom to verify more positions than a linear chain of 3 uses.

Tree drafting (EAGLE-2 / SpecInfer style) is the lever. Instead of one linear chain of K-1 proposals, propose a small tree, branching where the drafter is least confident, and verify the whole tree in one target forward with a tree attention mask. With verify at T=4 costing 39.95 ms, lifting emitted-per-round from 2.373 to about 3.2 would put the pairing near 2.4x. This is a substantially larger piece of work than Phases 0-4: it needs a tree mask on the verify path, a tree-aware walk, and tree-aware rollback.

Two honest caveats on comparing against vLLM and SGLang:

• Their headline speculative numbers generally come from CUDA deployments where decode is not purely memory-bandwidth-bound at B=1, so the arithmetic that makes a small drafter nearly free there does not transfer unchanged.
• EAGLE-3 style drafters used in those stacks ship a reduced draft vocabulary by design. Our drafter borrows a 248,320-wide head from the target, which is why Phase 2b is called out separately rather than folded into general cleanup.

Scope

In scope: Phases 0 through 4, on the shared MTP round loop and the Drafter trait, validated on the Qwen 3.8 27B pairing and regression-checked on Gemma 4.

Out of scope for this issue: tree drafting (file separately once Phases 0-4 land and the profile is re-measured), the verify-side early-exit walk (#1179), and anything that changes the target's numerics.

Blocker: MTP now declines on the hardware these numbers come from

This issue was filed before #1186 characterized the temperature-0 exactness failure and before #1189 turned the gate into a measured probe. With that probe in place, MTP fails closed on Apple GPU generation 15 and newer, which is exactly the hardware the entire profile below was measured on. So the work here cannot be validated through the default path until #1187 lands.

That does not invalidate the measurements. The round-cost split, the per-step drafter cost, and the M1 Ultra comparison were all taken through the production offline CLI before the gate change, and the drafter's cost has nothing to do with exactness: the drafter proposes, the target verifies, and no phase below alters the target's numerics.

It does change how the phases have to be run:

• Every measurement here needs the fix(speculative): gate MTP on a measured block-vs-chain exactness probe #1189 gate overridden (the opt-out that PR adds), and the override has to be named in the benchmark record so the numbers are not read as coming from a shipping configuration.
• Comparing temperature-0 output against classic decode is not a valid check on this hardware, because that comparison already fails for reasons unrelated to the drafter. Compare each phase against the same arm before the change instead.
• If perf(speculative): recover MTP throughput on Apple GPU generation 15+ without giving up temperature-0 byte-identity #1187 recovers exactness by forcing the plain qmv reduction on the verify path, the verify forward cost changes and the round split below shifts with it. Re-take the Phase 0 baseline after that lands rather than reusing these numbers.

Sequencing suggestion: land Phase 0 (the timers) regardless, since it is measurement-only and useful to #1187 as well. Hold Phases 1 through 4 until #1187 has settled what the verify side costs, or accept that their headline percentages will be restated afterwards.

Acceptance criteria

• Phase 0 lands first: per-component drafter timings plus an accept_hook_ms diagnostic, and the reported components sum to the measured draft_ms.
• Each subsequent phase reports a before and after on the same pairing and protocol used here, medians of 5 with cooldowns, load recorded.
• Temperature-0 output is unchanged by Phases 0, 1 and 2a, checked against the arm's own baseline rather than against classic decode. See the blocker note above for why the classic-decode comparison is not available on this hardware.
• Phases 2b and 3 report the acceptance-rate delta explicitly and do not become defaults without it.
• Gemma 4 MTP is re-benchmarked to confirm no regression on the shared paths.
• docs/benchmark_results/ gains an M5 Max record, and the M1 Ultra record is re-measured for the phases that land, since the drafter is the component that did not scale across generations.

References

• feat(inference): support the qwen3_5_mtp drafter for Qwen3.5/3.6/3.8 MTP speculative decoding #1165, feat(inference): support the qwen3_5_mtp drafter for Qwen 3.5 MTP speculative decoding #1182: the Qwen 3.5 MTP drafter and the M5 Max measurement this issue is derived from.
• fix(speculative): MTP temperature-0 byte-identity fails on Apple GPU generation 15+ and at block widths above the qmv batch limit #1186: why MTP declines on this hardware. The reason the measurements here need a gate override.
• perf(speculative): recover MTP throughput on Apple GPU generation 15+ without giving up temperature-0 byte-identity #1187: recovering exactness on generation 15+. Its step 3 points back here and asks that the two efforts be measured against the same round.
• fix(speculative): Gemma 4 MTP breaks byte-identity on generation 15+, and its gate returns unconditionally true #1188: the Gemma 4 arm of the same gate problem.
• fix(speculative): gate MTP on a measured block-vs-chain exactness probe #1189: the measured exactness probe that produces the decline.
• perf(inference): decide by measurement whether a device-side early-exit MTP verify walk beats the batched full-logits verifier #1179: verify-side device walk. Same class of problem, other side of the round.
• perf(inference): adaptive MTP enable/decline policy from per-pair profiling #333: adaptive MTP policy, which is where a confidence-gated dynamic block depth would eventually live.

[inureyes]

Phase 0 measured on M1 Ultra: PR #1194

Two results that bear on the phase ranking.

The accept hook is larger than draft_block. 14.1 ms per round against 10.3, on a Qwen3.8-27B pairing at block 3, 49 rounds, 129.1 ms per round. It had no counter so it sat in the residual, exactly as this issue says. It is a second drafter forward per round (append accepted tokens against the target's true hidden, precompute the next seed), and no phase currently addresses it.

Within the step, layers are 68.2% and the borrowed LM head is 19.7%; the host readback is 5.8%. That points at Phase 3 and Phase 2b ahead of Phase 1.

Full round split, unprofiled (the honest cost):

component	ms/round	share
verify forward (T=3)	100.5	77.8%
accept hook	14.1	10.9%
draft_block	10.3	7.9%
verify finalize	1.1	0.8%
residual	3.2	2.5%

A measurement problem had to be solved before any of that was reportable. MLX is lazily evaluated and a drafter step synchronizes exactly once, at the eval before the sampled id reaches the host. Timing the five stages as written puts 98% of the component time in readback_ms and leaves the rest measuring graph construction. So the split is opt-in behind MLXCEL_MTP_DRAFT_PROFILE=1, which costs 13.2% (11.83 ms/step against 10.45) because the syncs break pipelining, and the default mode reports dominant="unattributed" rather than naming a winner. The gap between the two totals bounds what Phase 1 can win by removing the round trip.

Per this issue's Blocker section: the run needed no MLXCEL_MTP_ALLOW_INEXACT override, because M1 Ultra is generation 13 and passes the exactness gate. These are shipping-configuration numbers. The M5 Max re-take will need the override and should say so in its record.

Benchmark record: docs/benchmark_results/mtp-drafter-step-profile-m1ultra-2026-08-17.md.

Caveat that matters for acting on this: the split is from an M1 Ultra and this issue's profile is from an M5 Max. The drafter step did not get faster across those generations (6.6 against 7.0 ms/step per this issue), which is suggestive but not evidence that the split transfers. Re-taking Phase 0 there is the next step, and it is the same command with the override set.

[inureyes]

Phase 0 on M5 Max: the decomposition transfers, and the absolute values say more than the shares

Second host for the Phase 0 measurement, on the hardware this issue is aimed at. Full write-up in #1197 (docs/benchmark_results/mtp-drafter-step-profile-m5max-2026-08-17.md).

Apple M5 Max, 40 GPU cores, 128 GB, macOS 26.6.1, Xcode 26.5, main at 33951285. Same target, drafter, prompt and command as the M1 Ultra run.

Blocker section, as asked

This is not a shipping configuration. The #1189 gate is fail-closed on generation 17, so the run required MLXCEL_MTP_ALLOW_INEXACT=1. The probe logged its own reading at startup:

MTP exactness probe FAILED but MLXCEL_MTP_ALLOW_INEXACT is set: verify block
position 0 differs from the single-token chain in 165506 of 496640 logit bytes.

33.3% of the logit bytes, against the 141k of 496k the M1 Ultra probe reports at block 12. Because MTP is inexact here the token stream differs from the M1 Ultra run (51 rounds at acceptance 0.6667 against 49 at 0.7320), so per-round and per-step figures compare cleanly and run totals do not.

The attribution transfers

component	M1 Ultra	M5 Max
layers	68.2%	73.5%
LM head (borrowed 248,320-wide)	19.7%	21.7%
readback	5.8%	3.0%
ids + embedding	6.0%	1.8%
slice + fused_sample	0.1%	0.02%

Same rank order, and the two largest gain share: together 95.2% of the profiled step here against 87.9% there. Profiling costs 0.62 ms/step (5.8%) here against 1.38 ms (13.2%) there, so quote 10.67 ms as the step cost and the profiled column only for attribution.

The finding is in the wall clock, not the shares

component	M1 Ultra ms/round	M5 Max ms/round	speedup
round total	129.1	58.5	2.21x
verify forward (T=3)	100.5	35.24	2.85x
accept hook	14.1	10.60	1.33x
draft_block	10.3	10.67	0.97x

The 27B verify forward is 2.85 times faster on this host. draft_block is not faster at all, and no component of it is faster per step (layers 8.06 to 8.29 ms/step, LM head 2.33 to 2.45). The drafter side goes from 18.8% of the round to 36.3% by standing still while everything around it accelerates.

That is this issue's premise measured rather than modeled, and it carries a corollary worth stating: the cost is not a property of the target model's speed, so it will not be absorbed by faster hardware. Why the step does not scale is not established here.

What this does to the phase ranking

The ranking is the same on both hosts, and the gap widens in the direction the M1 Ultra measurement already pointed.

• Phase 3 (4-bit drafter weights) targets the largest component, larger here at 73.5%.
• Phase 2b (reduced draft vocabulary) targets the second largest, also larger here at 21.7%.
• Phase 1 (remove the host round trip) is worth less on the faster host, not more. Readback is 3.0% against 5.8%, and the pipelining effect it relates to is 5.8% against 13.2%.

One thing that does not close

The non-step cross-check that agreed to 0.7% on M1 Ultra is off by 9.1% here: instrumented layers plus ids come to 508.43 ms against the same arm's accept_hook_ms of 559.58. The known set_seed_from_hidden gap (the hook's projection and readback sit inside accept_hook_ms without a split) points the right direction but does not explain an order-of-magnitude change across hosts. Instrumenting those two, already named as the next increment, would settle it. Reporting it rather than quoting the pair as agreement.