fix(speculative): compare three inputs before the exactness probe passes

docs/environment-variables.md

@@ -185,7 +185,7 @@ The OpenAI audio endpoints (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1
 | `MLXCEL_ENABLE_MTP_BATCH` | truthy value | off | **Advanced.** Forces the batched Gemma 4 MTP burst path for parity/debug testing. Not governed by the adaptive policy (issue #333), which scopes to the validated, byte-identical B=1 path. |
 | `MLXCEL_ENABLE_MTP_DEFERRED` | `1` | off | **Advanced.** Enables the deferred greedy verifier path for Gemma 4 MTP when sampling settings allow it. |
 | `MLXCEL_GDN_CHAIN_PARITY` | `0` to disable, any other value (or unset) to enable | on | **Advanced, diagnostic escape hatch.** Gates the chain-parity gated-delta Metal kernel used by Qwen 3.5 MTP's speculative verify and rollback-replay paths (issue #1165). The standard gated-delta kernel carries float32 recurrent state across a `T = K` verify block and rounds it to the storage dtype only once at the end, while the classic single-token decode chain rounds after every token; a `T = K` verify block is therefore NOT bit-identical to `K` consecutive single-token decode steps unless the state is rounded per in-block step. The chain-parity kernel (`gated_delta_step_seqpar`) does that rounding, which is what makes Qwen 3.5 MTP's temperature-0 output byte-identical to classic decode. **Setting this to `0` forfeits that exactness contract**, restoring the pre-#1165 block numerics for A/B attribution of the parity kernel's own cost and acceptance effect; do not set it to `0` in a deployment that needs byte-identical speculative output. Metal-only: the non-Metal ops fallback ignores the flag (the parity guarantee does not exist off Metal today). **The kernel is necessary but not sufficient**: byte-identity also requires every quantized projection to dispatch to the same MLX kernel at `M = block_size` as at `M = 1`, which is not true on every GPU generation or at every block width, so the runtime probe behind `MLXCEL_MTP_ALLOW_INEXACT` is what actually decides whether MTP engages. See `docs/benchmark_results/qwen38-mtp-m1ultra-2026-08-16.md` for the measured kernel cost (inside the dispatch-noise band). |
-| `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M >= 12` (27B operand sizes, Apple GPU generation 13/14, architecture size `d`) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12; the same checkpoint on an M5 Max diverges from block width 2. Set this to `1` when researching MTP throughput on hardware where the probe declines, accepting that the byte-identity contract no longer holds; the decline and the override are both logged at WARN with the differing-byte count. Read once per process. |
+| `MLXCEL_MTP_ALLOW_INEXACT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Engage Qwen 3.5 MTP speculative decoding even when the startup exactness probe reports that the multi-token verify block is **not** byte-identical to the single-token decode chain. Before enabling MTP the runtime now measures the property instead of predicting it: one synthetic verify block and the equivalent single-token chain are run from the same state on the loaded checkpoint at the configured `--draft-block-size`, and their logits are compared byte for byte (three independent synthetic inputs, each two short prefills plus `K + 1` forwards; measured 4.9 s for the first call and 1.3 s for a later one per input on a Qwen3.8-27B 4-bit target on an M1 Ultra, the difference being MLX's one-time kernel compilation; more than one input because a kernel pair can disagree by only a byte or two out of ten thousand, at which amplitude a single draw can read as equal; memoized per (model, block width) and warmed at worker startup so it never lands on the request path). A divergence means temperature-0 speculative output would silently differ from `mlxcel generate` without `--draft-model`, so the default is to decline and run classic decode. The static conditions (Metal backend, `supports_metal_gated_delta_kernel` geometry) still apply and are checked first; this probe covers what they cannot, namely which MLX kernel each **quantized projection** dispatches to at `M = K` versus `M = 1`. That choice depends on the GPU generation, the quantization mode, the operand sizes and the block width: `use_qmv_wide` in [`mlx/backend/metal/quantized.cpp`](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/quantized.cpp) sends `M >= 2` to a different reduction whenever `mode != "affine" || arch_gen >= 15`, and `get_qmv_batch_limit` sends `M` above 10, 12, 18 or 32 (by architecture size and generation) to the matrix-matrix kernel. Measured: an affine 4-bit Qwen3.8-27B target on an M1 Ultra is byte-identical at block widths 1 through 11 and diverges at 12 (the `arch_size == 'd'` branch); the same checkpoint on an M5 Max diverges from block width 2, where the `M >= 2` split fires before any batch limit can be observed. Within one checkpoint the limit is per projection, not per model: Gemma 4 12B's attention shapes hold to 17 on an M1 Ultra while its MLP shapes break at 12, so a model's own cliff is the minimum over its shapes, which is why this is measured rather than tabulated. Set this to `1` when researching MTP throughput on hardware where the probe declines, accepting that the byte-identity contract no longer holds; the decline and the override are both logged at WARN with the differing-byte count. Read once per process. |
 | `MLXCEL_MTP_TICK_SLICE` | `0`/`false`/`no`/`off` to disable, any other value (or unset) to enable | on | Tick-cooperative B=1 MTP serving (issue #734). When on (the default), a B=1 MTP request on the Gemma 4 family is served one speculative round per scheduler tick, alternating with the classic decode/prefill actions, so concurrent classic-decode rows advance between rounds and the head-of-line stall a speculative request imposes drops from the whole burst to about one round (`burst_wall_ms` in the finalize log reports the max single-tick wall). Tokens stream per round instead of in one end-of-burst lump. Set to an off value to restore the legacy run-to-completion burst (the whole request served inside one tick). The interleaving trades roughly 27% of the speculative request's own aggregate decode throughput (cross-tick round gaps) for that bounded stall, so a deployment serving speculative requests without concurrent classic traffic can turn it off to keep the full-throughput burst. Greedy output, acceptance accounting, and every other env gate are unchanged in both modes; DFlash and the batched B>1 paths always run to completion regardless of this flag. |
 | `MLXCEL_MTP_SLICE_GRANT_ROUNDS` | non-negative integer | `8` | Grant budget for one hold of the tick-slice speculative slot (issue #746), counted in executed slices (slice 0, the prefill + seed, counts as the first slice of a grant). While a slice is in flight, up to 2 further tick-slice-eligible requests park in a grant backlog instead of permanently falling back to classic decode; once the active request has run this many slices with the backlog non-empty, it parks at the next round boundary and the slot is granted to the next request (priority lane first, FIFO within a lane, with an anti-starvation floor: an entry passed over by 2 grant decisions is granted next regardless of lane), so concurrent long streams share speculative acceleration in bounded turns. The budget is read once per grant and per admission decision (cached for the per-round expiry check), so changes apply from the next grant. The budget binds only under contention: a single speculative request never rotates and behaves exactly as under #734. Rotation preserves per-request token streams byte-identically (the drafter is re-armed from the session's own stored verify output at every round). `0` disables rotation and restores the pre-#746 behavior: the active request holds the slot for its whole generation and every concurrent speculative request falls back to classic decode. Unparseable values fall back to the default. |
 | `MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT` | `1`/`true`/`yes`/`on` to enable; unset or anything else to disable | **off** | Acceptance-optimal speculative acceptance for the classic `SpeculativeGenerator` path (offline `mlxcel generate --draft-model`), issue #902. When on, `temperature > 0` verification uses modified rejection sampling (accept the drafted token `t` iff `u * q(t) <= p(t)` for a fresh `u ~ U[0,1)`, and on the first rejection emit a draw from the normalized residual `relu(p - q)`) instead of the default sampler-match rule (accept iff the draft equals an independent draw from the target sampler). **Both rules are distribution-preserving**: the emitted stream is a target-only sample either way, which is the central correction to the issue's premise. What changes is the acceptance probability, which rises from `sum_x p(x) q(x)` to `sum_x min(p(x), q(x))`, the maximal-coupling ceiling for any correct rule. **Opt-in rather than default** because the gain is the ratio between those two quantities and it collapses toward 1 whenever the drafter is confident (`q(t*) ~ 1` makes `min(p, q)` and `p * q` coincide); measured at about 1.02 on a Llama-3.1-8B / Llama-3.2-1B pair at temperature 0.7, which does not pay for two extra full-vocabulary passes and a host sync per verified position. Check the available gain with `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1` before enabling. Enabling changes the RNG stream, so at an equal seed the emitted tokens differ from a default run even though the distribution is identical. Greedy (`temperature == 0` or `top_k == 1`) never reaches either rule and is byte-identical. The Gemma 4 MTP and DFlash round loops are unaffected: they select the target token by argmax regardless of temperature, so this switch is inert there. `SpeculativeGenerator::with_stochastic_acceptance(bool)` overrides it programmatically. Read once per process. See [`speculative-acceptance.md`](speculative-acceptance.md). |

src/models/qwen3_5.rs

@@ -57,6 +57,17 @@ use std::path::Path;
 /// a fraction of a second on a 27B target.
 const PROBE_PROMPT_LEN: usize = 8;
 
+/// Independent synthetic inputs the exactness probe compares before it
+/// is allowed to report equality.
+///
+/// One is not enough. A kernel pair can differ by a byte or two out of
+/// ten thousand, and at that amplitude whether any byte differs depends
+/// on the draw, so a single input can report equality for arms that are
+/// on different kernels. Three, because the failure the extra draws
+/// guard against is a false *pass*, and each one costs about a second on
+/// a 27B target at worker startup.
+const PROBE_DRAWS: usize = 3;
+
 // Configuration.
 #[derive(Debug, Clone, Deserialize)]
 pub struct Qwen35Config {
@@ -1525,11 +1536,25 @@ impl Qwen35Model {
     /// for the reasons in
     /// [`crate::models::speculative_exactness`].
     ///
-    /// Cost: two 8-token prefills plus `block_size + 1` forwards on
-    /// throwaway caches. Measured on `models/qwen3.8-27b-4bit` on an M1
-    /// Ultra: 4.9 s for the first call in a process and 1.3 s for a later
-    /// one at a different width, the difference being MLX's one-time
-    /// kernel compilation for these shapes. Callers memoize through
+    /// Repeated over [`PROBE_DRAWS`] independent synthetic inputs, and
+    /// any single divergence is the answer. One draw is not enough: a
+    /// kernel pair can differ by only a byte or two out of ten thousand,
+    /// and at that amplitude whether *any* byte differs is itself
+    /// draw-dependent, so a single input can read `equal` for arms that
+    /// are genuinely on different kernels (measured at op level on
+    /// 2026-08-17: mxfp4 group 32 at 5120 -> 5120 moves 1 to 11 bytes of
+    /// 10240 depending on the draw, where the affine row at the same
+    /// shape moves 39%). Model-level amplification across 60-odd layers
+    /// makes a false `equal` far less likely here than at op level, but
+    /// "less likely" is not the property a correctness gate should rest
+    /// on when extra draws cost about a second each.
+    ///
+    /// Cost: per draw, two 8-token prefills plus `block_size + 1`
+    /// forwards on throwaway caches. Measured on
+    /// `models/qwen3.8-27b-4bit` on an M1 Ultra: 4.9 s for the first call
+    /// in a process and 1.3 s for a later one at a different width, the
+    /// difference being MLX's one-time kernel compilation for these
+    /// shapes. Callers memoize through
     /// [`crate::models::speculative_exactness::mtp_exactness_gate`] and
     /// the server warms it at worker startup, so it runs once per
     /// process and never on the request path.
@@ -1548,11 +1573,26 @@ impl Qwen35Model {
             return BlockChainExactness::NotRun("degenerate vocabulary");
         }
 
-        // Fixed synthetic ids. Values do not matter: MLX kernel dispatch
-        // is a function of shape, dtype, quantization mode and `M`, none
-        // of which vary with the tokens. Wrapped into the vocabulary so a
-        // tiny test config is still valid.
-        let wrap = |i: usize, stride: usize, offset: usize| ((i * stride + offset) % vocab) as i32;
+        for draw in 0..PROBE_DRAWS {
+            let verdict = self.probe_one_draw(block_size, vocab, draw);
+            if !verdict.is_equal() {
+                return verdict;
+            }
+        }
+        BlockChainExactness::Equal
+    }
+
+    /// One (prompt, block) draw of [`Self::probe_block_chain_exactness`].
+    fn probe_one_draw(&self, block_size: usize, vocab: usize, draw: usize) -> BlockChainExactness {
+        // Synthetic ids, varied per draw. Kernel *dispatch* does not
+        // depend on the values (it is a function of shape, dtype,
+        // quantization mode and `M`), but whether a small last-ulp
+        // difference lands on a byte that differs does, which is the
+        // whole reason for more than one draw. Wrapped into the
+        // vocabulary so a tiny test config stays valid.
+        let salt = draw * 977 + 1;
+        let wrap =
+            |i: usize, stride: usize, offset: usize| ((i * stride + offset + salt) % vocab) as i32;
         let prompt: Vec<i32> = (0..PROBE_PROMPT_LEN).map(|i| wrap(i, 7, 1)).collect();
         let block: Vec<i32> = (0..block_size).map(|i| wrap(i, 13, 3)).collect();
 

src/models/speculative_exactness.rs

@@ -30,13 +30,19 @@
 //!   is `mode != "affine" || arch_gen >= 15`. When it holds, `M >= 2`
 //!   takes `qmv_wide` while `M == 1` takes `qmv`, and the two reduce over
 //!   K with different lane counts. Measured: an affine 4-bit projection is
-//!   byte-equal at every `M` on an M1 Ultra (generation 13) and diverges
-//!   from `M = 2` on an M5 Max (generation 17+). An mxfp4 projection
-//!   diverges from `M = 2` on **both**.
+//!   byte-equal below the batch limit on an M1 Ultra (generation 13) and
+//!   diverges from `M = 2` on an M5 Max (generation 17+). An mxfp4
+//!   projection diverges from `M = 2` on **both**, by 1 to 11 bytes of
+//!   10240 rather than the 39% the affine row moves once it breaks.
 //! - `get_qmv_batch_limit` decides where the matrix-matrix kernel takes
-//!   over. On an M1 Ultra with both operands above 4096 that is 12, so a
-//!   block width of 12 or more forfeits byte-identity even there, and
-//!   `--draft-block-size` is user-settable.
+//!   over, and it is 10, 12, 18 or 32 depending on architecture size and
+//!   generation, not one number. With both operands above 4096 it is 12
+//!   on an M1 Ultra (`arch_size == 'd'`) and 10 on an M5 Max (the
+//!   `default` branch), so a wide enough block forfeits byte-identity
+//!   even on a generation whose `qmv` path is otherwise exact, and
+//!   `--draft-block-size` is user-settable. This threshold is only
+//!   observable where `use_qmv_wide` is false: from generation 15 the
+//!   `M >= 2` split fires first and hides it.
 //!
 //! Encoding those rules on our side would be correct today and would go
 //! stale the next time MLX retunes a threshold, which is the same failure
@@ -47,16 +53,30 @@
 //!
 //! ## What a passing probe does and does not prove
 //!
-//! The probe compares logits, not sampled tokens, and byte-equal logits
-//! imply an equal argmax at every position for *that* input. It runs one
-//! block on one synthetic prompt, so it does not prove equality for every
-//! future input directly. What it establishes is that the two arms take
-//! the same kernels: MLX dispatch is a function of shape, dtype,
-//! quantization mode and `M`, none of which vary with the token values.
-//! A byte-equal probe therefore means the block arm and the chain arm are
-//! running the same code on the same shapes, which is the mechanism the
-//! contract rests on. A diverging probe is conclusive in the other
-//! direction: the arms provably differ.
+//! A diverging probe is conclusive: the arms provably differ, and no
+//! amount of further sampling can make them agree.
+//!
+//! A passing probe is weaker than it first looks, and the earlier version
+//! of this note overstated it. It said that byte-equal logits mean both
+//! arms took the same kernels, since MLX dispatch depends on shape,
+//! dtype, quantization mode and `M` and not on token values. The
+//! dispatch half is true; the inference from it is not. Two *different*
+//! kernels can still produce byte-identical output on a particular input
+//! when their disagreement is small enough. Measured at op level on
+//! 2026-08-17: mxfp4 group 32 at 5120 -> 5120 moves 1 to 11 bytes of
+//! 10240 depending only on the operand draw, while the affine row at the
+//! same shape moves about 39%. In that low-amplitude regime a single
+//! input is a coin toss, and one M5 Max harness did read it as equal
+//! before a seed sweep showed otherwise.
+//!
+//! So the probe compares several independent inputs
+//! (`PROBE_DRAWS` in `models::qwen3_5`) and only reports equality when
+//! every one of them agrees. Model-level amplification helps too: a
+//! per-projection last-ulp difference propagates through 60-odd layers
+//! before it reaches the logits, so the amplitude at the point of
+//! comparison is far above the fragile regime. Neither argument is a
+//! proof for every future input; together they are what this gate rests
+//! on, and the failure they guard against is a false *pass*.
 //!
 //! ## Failure policy
 //!