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
Activity
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:
mode T=1 T=2 T=3 T=4 mxfp4 group 32 equal equal equal equal nvfp4 group 16 equal equal equal DIVERGES (1 / 10752) affine 4-bit group 64 (control) equal DIVERGES (501 / 10240) DIVERGES (501) DIVERGES (501) 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.
- added a commit that references this issue
on Aug 16, 2026 - added a commit that references this issue
on 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:host seeds result M5 Max 1165, 4242, 7, 99991 DIVERGES, 1 to 9 bytes of 10240 M1 Ultra 1165, 4242, 7, 99991, 31337, 555 DIVERGES, 1 to 11 bytes of 10240 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.
- added 4 commits that reference this issue
on Aug 16, 2026 - added a commit that references this issue
on Aug 20, 2026
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_targetinsrc/server/batch/speculative_burst.rs, which checksmetal_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_sizeversusM = 1, and that admits hosts where byte-identity does not hold.Evidence
Apple M5 Max 128 GB (macOS 26.6.1),
mainat115f6089block_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).' 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_chainpasses there, the kernel is provably dispatched (mutation-checked withMLXCEL_GDN_CHAIN_PARITY=0), and unset /=1/=0all 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 aT = Kcall againstKsingle-position calls at Qwen3.8-27B shapes.quantized_matmul5120 to 17408 (affine 4-bit, group 64)quantized_matmul5120 to 5120 (affine 4-bit)quantized_matmul5120 to 5120 (mxfp4, group 32)fast_rms_norm,fast_rope, causal SDPABlock 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_BLOCKoverride, 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 theqmvkernel.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 generatewithout--draft-model. On generation 13/14 the same happens for any--draft-block-sizeof 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, theMLXCEL_GDN_CHAIN_PARITYrow, anddocs/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=1as 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
MLXCEL_MTP_ALLOW_INEXACTis the only way to engage a diverging configuration, is off by default, and logs at WARN when used.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.