perf(inference): decide by measurement whether a device-side early-exit MTP verify walk beats the batched full-logits verifier

[inureyes]

Summary

The MTP verify pass gave up early-exit on the first draft mismatch because the only implementation available at the time cost K separate cxx/MLX round trips, which was worse than the tail projection it saved. A device-side walk would give the skip without the round trips. Decide by measurement whether it is worth building, and correct the stale comment that currently misdescribes why the existing path was chosen.

Applies to Gemma 4 MTP today, and to Qwen 3.5 MTP once #1165 lands.

What is actually in the tree, which is not what the comment says

src/models/gemma4_mtp_target.rs gates a use_deferred_greedy path behind MLXCEL_ENABLE_MTP_DEFERRED=1, temperature == 0, and logprobs disabled. The comment above that gate reads:

The upstream Python reference uses deferred greedy hidden->logits projection by default. In Rust/MLX today that path projects one position at a time across the cxx bridge and is slower than the batched [K, vocab] LM-head projection for Gemma 4 31B on local Apple Silicon runs.

That no longer describes the code. argmax_from_hidden_positions projects the whole block in one call:

let logits = self.wrapper.speculative_logits_from_hidden(hidden_full);   // all K positions
let argmax = mlxcel_core::argmax_last_axis(logits.as_ref().unwrap());    // device-side
mlxcel_core::eval(&argmax);                                              // one eval
materialize_argmax_i32_vec(&argmax, expected_len)                        // one host copy

speculative_logits_from_hidden(&self, hidden: &MlxArray) (src/models/gemma4.rs:4261) takes the whole tensor and does a single as_linear. There is no per-position projection variant anywhere in the tree (grep -rn "logits_from_hidden" src returns only these).

The function's own doc comment is current and explicit about the tradeoff that was actually made:

Greedy target-token extraction from pre-norm hidden states without a Rust-side per-position FFI loop. [...] It deliberately does not early-stop on the first mismatch: for the small Gemma 4 MTP block sizes we use today, avoiding K separate cxx/MLX calls is more important than skipping the tail projection on low-accept rounds.

git log -L shows both the gate comment and argmax_from_hidden_positions were last touched by fcf9e205 fix(perf): avoid slow Gemma 4 MTP singleton bursts. That commit rewrote the per-position loop into the one-graph form and dropped the early-stop, but left the gate comment describing the pre-rewrite behavior.

So the current state is:

• The "deferred" path defers nothing. It projects the same K positions as the default. The only difference is that it splits one bridge call (forward returning logits) into two (forward returning hidden with skip_final_norm, then a separate norm plus LM-head call).
• It is therefore strictly worse than the default: identical compute, one extra bridge crossing, one extra intermediate tensor. Marginally so, and for a different reason than the comment gives.
• The env flag advertises a capability the code does not provide.

This cost the orchestrator of the #1165 chain real time: the stale comment reads as "the escape hatch exists and lost on the merits", which is not what happened.

The actual opportunity

The early-exit was abandoned for an implementation reason, not because early-exit is worthless. At temperature == 0 the verify walk stops at the first position where the target argmax disagrees with the draft, so positions after that are computed and thrown away. A walk executed on device in one dispatch would skip them without paying K host round trips, dissolving the tradeoff the doc comment describes.

Prior: the arithmetic says this probably does not win at current block sizes

Worth stating up front so nobody builds it on hope.

The LM head for Qwen 3.5 / 3.8 is 5120 x 248320, about 1.27B parameters, roughly 636 MB at 4-bit. Gemma 4 is comparable. In a batched projection over K positions the weight matrix is read once and the operation is bandwidth-bound on that read. Skipping tail positions does not reduce it. What early-exit actually saves is the per-position output write and reduction: at K = 3, about 3 MB of logits versus 1 MB, plus two rows of argmax.

So the saving scales with K, and the current block sizes are small: Qwen 3.8's MTP config declares block_size: 3 (2 drafted tokens per round), and the Gemma comment references K = 4. At that size the prediction is that a device-side walk buys close to nothing, because the dominant term is invariant to it.

The condition under which it could win is a materially larger K. Note the Qwen MTP drafter sets prefer_requested_block_size = true (src/lib/mlxcel-core/src/drafter/mod.rs:662-663), so operators can request larger blocks and the drafter honors them rather than clamping to the configured depth. If large-block MTP turns out to be a real operating point, this becomes worth revisiting; if it does not, this issue closes as a measured no.

What to do

Per CLAUDE.md's performance-issue rules, this issue is "decide, by measurement at the real call site, whether the optimization belongs in the product", and a measured no closes it successfully.

1. Correct the stale comment first, independently of the rest. It currently misattributes the reason for the design and will keep misleading readers. One-line change, no measurement needed, worth doing even if everything below is declined.
2. Decide what to do with MLXCEL_ENABLE_MTP_DEFERRED. It gates a path that is strictly worse than the default and no longer does what its name says. Either remove it, or rename and re-document it as "split the verify forward from the LM-head projection", which is what it actually controls.
3. Sweep the block-size dependence before building anything. Measure the batched full-logits verifier at K of 3, 4, 8, 16, 32 on a real checkpoint, and separately measure the theoretical ceiling of early-exit at each K by measuring the projection cost for accepted+1 positions versus K positions. If the gap is inside measurement noise at the K values anyone actually runs, stop here and close with the numbers.
4. Only if the sweep shows a real gap, implement the device-side walk. Both arms must be measured through the production verify path, not as a microbenchmark: op-level wins in this tree have repeatedly evaporated at the real call site (see the Epic: Fused paged decode, sorting-free sampling, and serving performance techniques #909 epic notes in CLAUDE.md).
5. Whatever the outcome, record the numbers in docs/benchmark_results/<feature>-<hw>-<date>.md and link it here.

Guard the measurement

The two arms must be provably different. A harness that cannot tell them apart will report noise as a result, which happened three separate times in epic #909. Require a counter or log line that can only move in one arm, for example the count of positions actually projected per verify round, and make the harness refuse to print a table when that attribution fails. Follow the guard pattern in examples/sparse_paged_decode_bench.rs.

Report medians and dispersion over repeated runs and record the machine load. This development box carries background load, so within-run ratios survive it and absolute timings do not.

Out of scope

• Changing the default verify path before a measurement justifies it.
• Temperature above 0 and logprobs-enabled requests. The existing gate already restricts the deferred path to greedy without logprobs, and early-exit is only meaningful for greedy verification.
• Batched (B > 1) MTP.

Acceptance criteria

• The use_deferred_greedy gate comment describes the code that exists
• MLXCEL_ENABLE_MTP_DEFERRED either removed or renamed and re-documented to match what it controls
• Block-size sweep measured on a real checkpoint through the production verify path, with the projected-position counter proving the two arms differ
• A decision recorded with numbers: build the device-side walk, or close as a measured no
• If built: measured through the production path, medians and dispersion, machine load recorded, default set by the measurement rather than by expectation
• Benchmark record written to docs/benchmark_results/ and linked from this issue

[inureyes]

Split out of #1165 during that issue's implementation, after the orchestrator found that the MLXCEL_ENABLE_MTP_DEFERRED path does not actually defer anything and that the comment explaining it describes a pre-fcf9e205 implementation. The comment fix in item 1 is worth doing on its own even if the measurement work is never scheduled.

[inureyes]

Measured: not built. PR #1296 carries the cleanup and the record.

The sweep this issue asked for is in docs/benchmark_results/mtp-verify-early-exit-decision-m5max-2026-08-22.md (PR #1296), taken with a new guarded harness (examples/mtp_projection_width_bench, each row asserts its [1, W, vocab] logits shape) plus a production block-size sweep on both pairings via the #1295 width-pin mode.

The prior in this issue was wrong in the letter and right in the verdict. The projection cost is not flat in W: W = 1 runs at the bandwidth roof, the M >= 2 kernels sag to ~126-142 GB/s by W = 8 under the production narrow selection, and the matrix-matrix kernel makes W = 16 cheaper than W = 8. So the early-exit ceiling at production widths is 2 to 4% of the verify forward rather than zero, and the walk still cannot collect it: a sequential walk re-reads the full 566-715 MB head per walked position, costing 1.2 to 1.6x the batched projection at the measured operating points (K = 3-4, acceptance 0.80-0.88). At K = 8 the walk would net ~1.9 ms, but K = 8 is not an operating point on either pairing and padding the projection past the qmv batch limit collects the same saving with a reshape (2.56 ms vs 4.51 ms at width 8). That pad-to-qmm observation is the lever to measure first if wide blocks ever become real.

Items 1 and 2 are also done in the PR: the stale gate comment is replaced by the measured reason, and MLXCEL_ENABLE_MTP_DEFERRED is removed with its orphaned helper chain (it split the verify forward from the projection without deferring anything). Byte identity re-verified after the removal.