Repository navigation
fix(speculative): MTP temperature-0 byte-identity fails on Apple GPU generation 15+ and at block widths above the qmv batch limit #1186
Description
Problem / Background
PR #1182 (issue #1165) shipped Qwen 3.5 MTP speculative decoding with a temperature-0 contract: the multi-token verify block emits the same tokens as the single-token decode chain. The gate protecting that contract is mtp_capable_target in src/server/batch/speculative_burst.rs, which checks metal_is_available() && supports_chain_parity_kernel().
Both conditions are necessary and neither is sufficient. The contract also depends on which MLX kernel each quantized projection dispatches to at M = block_size versus M = 1, and that admits hosts where byte-identity does not hold.
Evidence
Apple M5 Max 128 GB (macOS 26.6.1), main at 115f6089
block_verify_chain_matches_single_token_chainFAILS at the shipped block width 3: 2 diverging positions over 241 tokens, first at chain index 48 (single-token argmax 11592, block argmax 5587).adapter_chain_with_perfect_drafts_matches_classic_chainandadapter_chain_with_rejected_drafts_matches_classic_chainalso fail (1 position each, chain index 151).- The production server path shows it too: same prompt, temperature 0, 150 tokens, one server per arm, token ids compared through logprobs. Token ids agree for 13 tokens, then diverge at index 13 (classic
' broken', MTP' illustrated').
It is not the gated-delta chain-parity kernel. The synthetic bitwise pin parity_kernel_block_is_bitwise_equal_to_single_token_chain passes there, the kernel is provably dispatched (mutation-checked with MLXCEL_GDN_CHAIN_PARITY=0), and unset / =1 / =0 all produce the identical failure. It is also not the M5 neural-accelerator attention route, which was A/B'd with that path disabled.
Apple M1 Ultra 128 GB (macOS 26.6.1), same commit
A new model-free per-op diagnostic (tests/metal_block_vs_chain_op_parity.rs) compares a T = K call against K single-position calls at Qwen3.8-27B shapes.
| op (T = 3) | M5 Max | M1 Ultra |
|---|---|---|
quantized_matmul 5120 to 17408 (affine 4-bit, group 64) |
DIVERGES | equal |
quantized_matmul 5120 to 5120 (affine 4-bit) |
DIVERGES | equal |
quantized_matmul 5120 to 5120 (mxfp4, group 32) |
not run | DIVERGES (3, 4, 2 of 10240 bytes) |
| dense matmul f16 5120 x 5120 | DIVERGES | DIVERGES (27, 27, 33 of 10240) |
| dense matmul f32 5120 x 5120 | DIVERGES | DIVERGES (5486, 5451, 5514 of 20480) |
fast_rms_norm, fast_rope, causal SDPA |
equal | equal |
Block width sweep on the M1 Ultra, both affine quantized shapes: equal at 1 through 11, DIVERGES from 12 onward.
Mechanism, from the MLX source
In https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp:
use_qmv_wide(mode, d)returnsmode != "affine" || d.get_architecture_gen() >= 15. Insidedispatch_qmv,M >= 2 && use_qmv_wide(...)routes toqmv_wide, otherwiseqmv. The two reduce over K with different lane counts (qmvusesbk = 32withbn = 8output rows; affineqmv_wideusesk_lanes = 8and streams up to 5 input vectors per weight tile), so their last-ulp results differ by design. That explains why generation 13 (M1) is byte-equal for affine and generation 17+ (M5) is not fromM = 2, and why mxfp4 diverges on every generation.get_qmv_batch_limit(K, N)returns 12 for both operands above 4096 on generation 13/14 with architecture sized, andM >= vector_limittakes the matrix-matrix kernel. That is the measured M1 Ultra cliff at 12.
Confirmation
With a new MLXCEL_TEST_CHAIN_BLOCK override, the M1 Ultra at block width 12 reproduces the M5 Max block-3 failure exactly: 2 diverging positions over 241 tokens, first at chain index 48, single-token argmax 11592, block argmax 5587. Same count, same index, same token id pair. The single variable is whether the verify block leaves the qmv kernel.
Impact
On Apple GPU generation 15 and newer (M3, M4, M5), a user running Qwen 3.5 MTP at temperature 0 gets output that silently differs from mlxcel generate without --draft-model. On generation 13/14 the same happens for any --draft-block-size of 12 or more, and for any non-affine (mxfp4/nvfp4) checkpoint.
The docs state the byte-identity contract unconditionally for the Metal path (docs/environment-variables.md, the MLXCEL_GDN_CHAIN_PARITY row, and docs/benchmarks.md).
Proposed Solution
Replace the static gate with a runtime probe. At MTP enable time, run one synthetic verify block and the equivalent single-token chain from the same state on the loaded checkpoint at the configured block width, and compare the logits byte for byte. Fail closed on divergence, with MLXCEL_MTP_ALLOW_INEXACT=1 as an explicit, loud opt-in for throughput research.
Predicting the answer from a hardware table would be correct today and go stale the next time MLX retunes a threshold, which is the same failure mode as the gate being replaced.
Acceptance Criteria
- The gate measures byte-identity on the loaded checkpoint at the configured block width instead of inferring it from hardware and geometry.
- The measurement is memoized per (model, block width) and warmed at worker startup so it never runs on the request path.
- A diverging probe declines MTP and falls back to classic decode on both the server and the offline CLI paths, with the reason logged including the differing byte count.
-
MLXCEL_MTP_ALLOW_INEXACTis the only way to engage a diverging configuration, is off by default, and logs at WARN when used. - Documentation states the contract conditionally and names the measured boundaries.
- A test asserts the safety direction: the probe never reports equal where the real block-vs-chain comparison diverges.
Out of Scope
The Gemma 4 MTP family takes the same quantized projections and its gate arm still returns unconditionally true. It is exposed to the same generation 15+ behavior and needs the same probe, tracked in #1188.
Activity
inureyes commentedon Aug 16, 2026
Two narrowing corrections from M5 Max (generation 17, architecture size Max), measured at the same shapes this issue uses. Details and the full tables are on #1188.
1.
get_qmv_batch_limitis 10 here, not 12. The 12 in this issue is thearch_size == 'd'branch, which is the Ultra. A Max falls into thedefaultbranch, which returns 10 for both operands above 4096 on a generation that is not 13 or 14. Visible directly in a magnitude sweep at 5376 -> 5376: the worst differing-byte count steps atT = 2(539 / 10752,use_qmv_wide) and again atT = 10(1693 / 10752, the batch limit). Both thresholds show up on one machine, so the cliff should be stated as "get_qmv_batch_limit(K, N), which is 10 or 12 or up to 32 depending on architecture size and generation" rather than as a fixed 12.2. "mxfp4 diverges on every generation" does not hold on generation 17. Measured at 5120 -> 5120, the exact shape this issue used for the M1 Ultra mxfp4 row, with the affine case in the same run as a control:
The affine control breaking in the same process is what makes this a measurement rather than a broken harness. So the two modes invert across the two generations: M1 Ultra is affine-clean and mxfp4-broken, M5 Max is affine-broken and mxfp4-clean.
That inversion is not what
use_qmv_wide(mode, d) = mode != "affine" || arch_gen >= 15predicts on its own, so there is at least one more dispatch condition in play for the fp modes that neither of us has traced. I did not chase it, because it does not change what this issue asks for.It does sharpen the argument for #1189 over a hardware table, though: two generations times two quantization modes, and the answer flips in both directions. Any static table would have to encode the fp-mode behavior correctly on hardware nobody has measured yet.
Suggested edit to the Impact section: keep the generation 15+ affine statement and the batch-limit statement, and replace "and for any non-affine (mxfp4/nvfp4) checkpoint" with something measured, since it is confirmed on generation 13 and refuted on generation 17.
fix(speculative): gate MTP on a measured block-vs-chain exactness probe
test(speculative): make the op-parity operand seed overridable
inureyes commentedon Aug 16, 2026
Retraction: the mxfp4 half of the narrowing above does not hold
The 2026-08-16 comment on this issue made two corrections. The first, that
get_qmv_batch_limitis 10 on an M5 Max rather than the 12 this issue quoted, stands and has been folded into the in-tree documentation. The second does not.That comment reported mxfp4 group 32 at 5120 -> 5120 as
equalatT = 1through4on generation 17, and used it to retract "mxfp4 diverges on every generation". Re-measured through the in-tree diagnostic (tests/metal_block_vs_chain_op_parity.rs) on 2026-08-17, that shape diverges on both machines and on every seed tried:So the original statement was right and the narrowing was wrong. The harness behind the 08-16 reading (
tests/m5_block_vs_chain_op_parity.rs) lived on the closed #1190 branch and cannot be re-run; the reproducible measurement is the in-tree one, so that is the one that stands.Why one measurement read it as clean
Amplitude. At this shape the affine 4-bit row moves about 39% of the output bytes once it breaks, and the mxfp4 row moves 0.01% to 0.09%. In that regime whether any byte differs on a given draw is itself draw-dependent, so a single operand draw can genuinely produce zero differing bytes for arms that are on different kernels. It is not a sampling error in the statistical sense; the comparison is exact for the input it was given. The input was the problem.
MLXCEL_TEST_OP_PARITY_SEEDnow makes that check one command rather than a source edit.This changes the probe that closed this issue
The gate landed in #1189 compares one synthetic input per (model, block width) and the module documentation argued that a byte-equal result proves both arms took the same kernels, since MLX dispatch depends on shape, dtype, mode and
Mrather than on values. The dispatch half is true. The inference from it is not: this measurement is a direct counterexample, where different kernels produced byte-identical output on some draws.Model-level amplification makes a false pass much less likely there than at op level, because a per-projection last-ulp difference propagates through 60-odd layers before reaching the logits. But "much less likely" is not what a correctness gate should rest on when the fix costs about a second. Follow-up in flight: the probe now compares three independent synthetic inputs and reports equality only when all of them agree, and the overstated claim in the module documentation is replaced with what the evidence actually supports.
Filed here rather than as a new issue because it is a correction to this issue's own record and to the fix that closed it.
test(speculative): extend the block-vs-chain diagnostic to Gemma 4 sh…
fix(speculative): compare three inputs before the exactness probe passes
fix(speculative): compare three inputs before the exactness probe pas…
feat(speculative): let the MTP gate buy back exactness by dropping qm…
fix(speculative): gate MTP on a measured block-vs-chain exactness pro…
fix(speculative): gate MTP on a measured block-vs-chain exactness pro…
test(speculative): extend the block-vs-chain diagnostic to Gemma 4 sh…
fix(speculative): compare three inputs before the exactness probe pas…
feat(speculative): let the MTP gate buy back exactness by dropping qm…