fix(speculative): probe Gemma 4 MTP exactness instead of assuming it - #1258
Merged
Conversation
The Gemma 4 arms of mtp_capable_target returned true unconditionally, advertising a temperature-0 byte-identity the hardware does not always provide: on Apple GPU generation 15+ the default qmv_wide dispatch makes the verify block diverge from the single-token chain systematically (measured on #1188), and the 2026-08-19 three-host sweep caught a real prose divergence slipping through the unprobed gate on M1 Ultra as well. Gemma4Wrapper now carries the same block-vs-chain probe the Qwen 3.5 family runs: three synthetic draws on throwaway caches built directly from the inner model (touching neither the wrapper's internal slot nor any scheduler-owned seq_id slot), with both arms projected through the tied LM head at their own width so the head's M = K versus M = 1 dispatch is covered too. The verdict routes through mtp_exactness_gate, so a failing probe first retries with qmv_wide disabled and keeps it off when that restores exactness (~23% on this family's verify forward), and declines to classic decode otherwise unless MLXCEL_MTP_ALLOW_INEXACT is set. Wired in both consumers: the server burst dispatch and the offline CLI gate, which gets a Gemma-specific decline message. docs/benchmarks.md drops the 'not probed' qualifications: pre-gate rows are labelled as fast-kernel measurements, and the M1 Ultra prose case is called out as the known test of the probe on a host where qmv_wide is not the mechanism. Refs #1188
inureyes
added
status:review
This was referenced Aug 20, 2026
inureyes
added a commit
that referenced
this pull request
Aug 22, 2026
## Summary Issue #1261 gated any scoping of the qmv_wide narrow pin on first measuring what the rest of the process pays for it, with an explicit exit condition: if the tax is small, document the number and stop. Measured on the required generation 15 host (M3 Ultra, 512 GB): the B = 2/4/8 batched-decode tax is at most 1% on Gemma 4 31B and at most 0.2% on Qwen 3.8 27B (8 and 4 samples per arm per cell, every spread at or under 1%, both arms pinned explicitly via `MLXCEL_QMV_WIDE`). The reason is structural rather than numerical: both families the pin can fire on decode batches as per-sequence M = 1 forwards (Gemma 4 dispatches one row per forward through `forward_with_sequence_id`: on the VLM route the measured checkpoint takes, via `Gemma4VLModel`'s `forward_batched_with_context_and_ids` override and `forward_batched_with_seq_ids_dispatch`; on the text-only route, via the trait default's per-row loop. Qwen 3.5's override takes its per-row branch for every single-token input), so batched decode never dispatches the kernel the pin disables. The collateral cost that is real is one prompt-cache-adopted suffix prefill per request, whose M lands in the qmv window: +15.4 ms per forward on Gemma (+33% on cache-hit TTFT), +12.6 ms on Qwen (+28%), each fitting a single per-forward constant across all concurrency levels, with uncached and chunked prefill measured as controls (+0.2 to +1.1%). The mixed arm (one MTP stream holding the speculative slot plus four classic streams, Qwen pairing) bounds the bystander cost at 0.1 to 2.6%, and the loss tracks the MTP stream's own narrow verify occupying the shared worker, not the classic streams' kernels. Per the issue's exit condition, Step 2 is documented as not built; the boundary that would reopen it (a real joint batched decode for an MTP family, as Gemma 3 and Llama 4 already have) is named in the record. The measurement also settled the stale `MLXCEL_MTP_ALLOW_INEXACT` recipe. Since #1199 the gate retries without qmv_wide before the override is consulted, so on generation 15+ hosts where the narrow retry passes the override alone is inert: verified live with all four env recipes on the pairing #1258 measured, with the log lines, byte-identical output between the default env and the override-alone run, and the 117 vs 139 tok/s throughput split as evidence. The fast-kernel recipe is `MLXCEL_QMV_WIDE=1 MLXCEL_MTP_ALLOW_INEXACT=1`. Note the issue body's own description of its mixed arm ("`MLXCEL_MTP_ALLOW_INEXACT=1` with the switch left wide") needs the same correction; the measured arm B used the two-variable recipe. The docs are corrected accordingly, and the merged PR descriptions of #1199 and #1258 carry a prepended dated note pointing here. Found on the way and recorded for follow-up: on current `main` the 31B + bf16 assistant pairing probes non-identical under BOTH kernels on this host ("Disabling qmv_wide did not make it exact either"), so the default-env gate declines the batch-capable burst #1217 enabled; #1217's 1.95x to 2.65x rows were measured at `9e2c6675`, which predates #1258's Gemma probe. ## What changed - `scripts/bench_qmv_wide_pin.sh` (new): ABBA boot driver for both measurement arms (`sweep`: pinned wide/narrow B-sweep without a drafter; `mixed`: one MTP stream plus N classic streams, with each boot's arm identity grepped from the gate log lines). - `scripts/bench_qmv_pin_mixed.py` (new): mixed-workload client; classic decode rates are the reported quantity, and a window is invalid unless the MTP stream decoded through at least 95% of it. - `scripts/bench_serving_concurrency.py`: count `reasoning_content` deltas alongside `content`; a reasoning model that spends its budget thinking previously reported no TTFT and no decode rate at all. - `docs/benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md` (new): the full measurement record, including the contaminated long-context cell reported as text-divergence-contaminated rather than averaged in. - `docs/benchmarks.md`: link the record from the pin discussion; correct the fast-row reproduction recipe and the "declining probe" sentence for the post-#1199 retry ordering. - `docs/environment-variables.md`: correct the `MLXCEL_MTP_ALLOW_INEXACT` row and add the missing `MLXCEL_QMV_WIDE` row. - `docs/benchmark_results/mtp-drafter-step-profile-m5max-2026-08-17.md`: dated note that its reproduction recipe predates the #1199 retry. - `.gitignore`: ignore the raw `bench-results/` run directories the driver writes. - `scripts/with_indexers_paused.sh`: refuse to suspend anything without a resume list, validate `INDEXER_RESUME_DEADLINE` as an integer instead of splicing it into a detached `bash -c` string, and record each pid before it is stopped. - `TECHNICAL_REPORTS/1278-qmv-wide-pin-tax-measurement-20260822.{en,ko}.md` (new): the bilingual technical report for this measurement. ## Changes during review The measurement and its conclusion are unchanged; the review cycles hardened the harnesses that produced it and corrected one code citation. - The Gemma 4 mechanism originally cited the text-only trait default. The measured checkpoint carries `embed_vision.*` weights, so it routes through `Gemma4VLModel`, which does override `forward_batched_with_context_and_ids`. That override delegates to `forward_batched_with_seq_ids_dispatch`, itself an explicit per-row loop, so the per-sequence `M = 1` conclusion holds on the route the run actually took. The Summary above cites both routes. - The driver could not be interrupted: a `trap` handler that returned resumed the loop, so an interrupt mid-run killed one server and booted the next arm, completing all boots and exiting 0. Handlers now exit, and `HUP` was added so a dropped connection cannot orphan a 31B server. - `wait_ready` probed the health endpoint before checking that the launched child was alive, so a server already holding the port would answer and every boot would measure one kernel while reporting success. Arm identity is the only property the ABBA design establishes, so this is now a port preflight followed by a liveness check. - An inherited `MLXCEL_*` export could silently collapse the two arms of the mixed run. It is now rejected up front, and the run records the inherited environment. - A `TypeError` from a malformed `usage` field was outside the surrounding except tuple in `bench_serving_concurrency.py` and propagated through the gather, killing every remaining concurrency level. It is now handled per level. - Harness exit status was discarded in both modes, so a failed window could unbalance the design silently. Failures are now recorded and surfaced. ## Test plan - [x] B-sweep, Gemma 4 31B: 8 boots ABBA, 1 warm-up + 2 measured passes each, under `scripts/with_indexers_paused.sh`, Time Machine off; per-stream generation lengths verified equal across arms (59 tokens) via `usage` counts. - [x] B-sweep confirmation, Qwen 3.8 27B: 4 boots ABBA, same protocol. - [x] Mixed arm, Qwen pairing: 4 boots A B B A, 1 warm-up + 2 measured windows each; arm identity verified from the exactness-gate log line of every boot; window validity (MTP overlap at least 99.7%) checked. - [x] Gate recipes: 4 env recipes, 2 interleaved samples each, on `gemma-4-12b-it-4bit` + assistant; log lines, byte-identity diffs, and throughput recorded. - [x] `python3 -m py_compile` on both Python harnesses; `bash -n` on the driver. No Rust source changed. - [x] Post-review: `bash -n` on both shell scripts, `python3 -m py_compile` on both Python harnesses, and `cargo fmt --check` all clean. Interrupt handling, port preflight, and the mktemp failure path each verified against a replica or stub. Closes #1261
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The Gemma 4 arms of
mtp_capable_targetreturnedtrueunconditionally, advertising a temperature-0 byte-identity the hardware does not always provide: #1188 measured a systematic divergence from classic decode on Apple GPU generation 15+ under the defaultqmv_widedispatch, and the 2026-08-19 three-host sweep caught a real prose divergence slipping through the unprobed gate on M1 Ultra as well.Gemma4Wrappernow carries the same block-vs-chain probe the Qwen 3.5 family runs, and both consumers route through it: the server burst dispatch (mtp_capable_target) and the offline CLI gate, which gets a Gemma-specific decline message. On a failing probe the shared gate retries withqmv_widedisabled and keeps it off when that restores exactness;MLXCEL_MTP_ALLOW_INEXACT=1remains the loud opt-out for the fast kernel.Related issues
Closes #1188.
Type of change
fix— the gate advertised a contract it did not checkProbe design
internalslot nor any scheduler-ownedseq_idslot is touched. That sequencing constraint is what kept these arms on an unconditionaltruebefore.forward_with_caches_and_embeddings), so the probe covers the head'sM = KversusM = 1dispatch as well as the decoder layers. The head shape (3840 -> 262144) was never measured in the test(speculative): extend the block-vs-chain diagnostic to Gemma 4 shapes #1192 op sweep, and covering it matters for the M1 Ultra case below.mtp_exactness_gate; nothing in that machinery changed except the retry log line, which now quotes the measured cost for both families.The default this picks, and why
#1188 left the failing-probe policy open: pay ~23% for byte-identity, or keep the fast kernel and say so. This PR picks consistency with the shipped Qwen behavior (exactness by default,
MLXCEL_MTP_ALLOW_INEXACTto opt out), for the reason the probe machinery's own docs give: a silent loss of byte-identity is worse than a lost speedup, and the opt-out is loud and reversible per process. The issue's follow-up measurements support treating this as the default rather than a coin flip: a conditional guard is not available (the kernel error and the top-two gap overlap; measured break-even fails), so the two options really are the endpoints, and one of them breaks an advertised contract silently.Measurements (Apple M5 Max 128 GB, macOS 26.6.1,
ef562bae)Probe behavior: on this generation-17 host the probe fails under
qmv_wide(verify position 0 differs in 275832 of 524288 logit bytes) and passes without it, so the gate dropsqmv_wideand engages MTP. One log line says what was traded.Byte identity (
gemma-4-12b-it-4bit+ 4-bit assistant, block 5, 250 tokens, temperature 0, prompt echo and banners stripped per #1188's methodology note):MLXCEL_QMV_WIDE=0The prose row's second cell is the prefill kernel effect #1188 documented: classic's own multi-token prefill takes
qmv_wideunder the default env, soclassic(wide on)andclassic(wide off)already differ from each other on that prompt. The contract the probe enforces, block equals chain under one kernel selection, holds on both prompts.Throughput (code prompt, 300 tokens, indexers paused, Time Machine stopped, 2 samples per arm):
MLXCEL_MTP_ALLOW_INEXACT=1(fast kernel)This reproduces #1188's decision numbers (43.5 / 93.19 / 121.50): the default now ships the honest 2.13x instead of advertising 2.76x with an unchecked contract, and the fast row is one env var away.
What this does not settle
The M1 Ultra case. The 2026-08-19 sweep measured a prose divergence there even though generation 13 never takes
qmv_wide, so the mechanism is different (the untested 262k LM head dispatch is the candidate this probe now covers). What the probe reads on that host, and whether the prose row still diverges behind a pass, needs a run on that machine;docs/benchmarks.mdnow names it as the known test of the probe.Docs
docs/benchmarks.mddrops the three "Gemma 4 is not probed" qualifications: pre-gate rows are labelled as fast-kernel measurements, reproducing them needs the opt-out, and the M1 Ultra case is recorded as above.Test plan
cargo fmt --checkcargo clippy --workspace --all-targets --features metal,accelerate -- -D warningscargo test --release -p mlxcel --liband-p mlxcel-core: failure sets identical tomainon this machine (the pre-existing local numeric reds; verified by diffing sorted failure lists against a stashedmainrun)