fix(speculative): probe Gemma 4 MTP exactness instead of assuming it

docs/benchmarks.md

@@ -222,12 +222,14 @@ by diffing the two completions. The probe is not a formality: whether a `T = K`
 verify block is bit-equal to `K` single-token steps depends on which MLX kernel
 each quantized projection dispatches to at `M = K` versus `M = 1`, which varies
 by Apple GPU generation, quantization mode and block width. The Qwen 3.5 family
-declines to classic decode when the probe fails (#1186). The Gemma 4 arms are
-not probed at all (#1188), and diffing them settles the question in neither
-direction: on M3 Ultra and on M5 Max two of the three prompts below diverge
-from classic decode and the third is byte-identical, so the claim holds per
-prompt rather than per pairing, and which prompts fall on which side tracks
-acceptance rather than the hardware. M4 is still unmeasured.
+declines to classic decode when the probe fails (#1186), and since #1188 the
+Gemma 4 arms run the same probe: on a failing probe the gate first retries with
+`qmv_wide` disabled and keeps it off when that restores exactness (about 23% on
+this family's verify forward), and declines otherwise. Gemma 4 rows measured
+before that gate landed are the fast kernel, not the byte-identical one; the
+row's record says which. One caveat the probe inherits: a passing probe is
+measured evidence, not proof, and the M1 Ultra prose-prompt divergence recorded
+on 2026-08-19 (three-host sweep) is the known case to re-test against it.
 
 For each pairing, record both the baseline (no drafter) and the MTP run:
 
@@ -644,8 +646,11 @@ at all, so B=1 is also its only decode path. The batch-capable 31B + bf16
 assistant measures ~1.2 to 1.4x on M5 Max. Set `MLXCEL_ENABLE_MTP_B1=0` to
 opt out on hardware where the B=1 verify forward does not pay for itself.
 
-Gemma 4 is not probed yet (#1188), so the rows above are the fast kernel rather
-than the byte-identical one. Keeping byte-identity on the code row, by dropping
+The Gemma 4 rows above were measured before the #1188 gate landed, so they are
+the fast kernel rather than the byte-identical one; with the gate in place the
+default on generation 15+ is the byte-identical kernel, and reproducing the
+fast rows needs `MLXCEL_MTP_ALLOW_INEXACT=1`. Keeping byte-identity on the
+code row, by dropping
 `qmv_wide`, measures 93.2 tok/s instead of 121.0 on M5 Max, or 2.14x instead of
 2.79x, and 117.5 tok/s instead of 138.5 on M3 Ultra, 1.83x instead of 2.16x.
 That is 23% of throughput on one host and 15% on the other, which is not the
@@ -655,8 +660,12 @@ where the drafter step and the accepted-token emission are unaffected. Quote
 whichever one the question is about, and not the other.
 
 On M1 Ultra there is no such cost, because generation 13 never takes
-`qmv_wide` in the first place, but the missing probe still lets a divergence
-through. The two arms have now been diffed at `temperature 0` on the three
+`qmv_wide` in the first place, but a divergence still got through while the
+arms were unprobed, and it is the case that tests the probe now that #1188
+routes these arms through it: the mechanism there cannot be `qmv_wide`, so
+what the probe reads on that host (and whether the prose row still diverges
+behind a pass) needs its own run. The two arms have now been diffed at
+`temperature 0` on the three
 prompts above on one host from each of three GPU generations, with the probed
 Qwen pairing run beside them as the control:
 

src/commands/generate.rs

@@ -1709,6 +1709,19 @@ where
 ///
 /// A target that is not MTP-capable returns a clear error instead of silently
 /// falling back, matching the issue's contract.
+/// Decline message for a Gemma 4 target whose exactness probe failed even
+/// with the `qmv_wide` retry (issue #1188). Mirrors the Qwen 3.5 message.
+fn gemma4_mtp_declined(block_size: usize) -> anyhow::Error {
+    anyhow!(
+        "Gemma 4 MTP speculative decoding declined: at --draft-block-size \
+         {block_size} this GPU's multi-token verify block is not byte-identical \
+         to the single-token decode chain, so temperature-0 output would \
+         silently differ from `mlxcel generate` without --draft-model. Try a \
+         smaller --draft-block-size, or set MLXCEL_MTP_ALLOW_INEXACT=1 to \
+         engage anyway and forfeit the byte-identity contract."
+    )
+}
+
 fn run_offline_mtp(
     model: &mlxcel::LoadedModel,
     draft_model_path: &Path,
@@ -1763,6 +1776,29 @@ fn run_offline_mtp(
         ));
     }
 
+    // Gemma 4 exactness gate: same call as the server's `mtp_capable_target`,
+    // for the same reason as the Qwen block above. These arms used to admit
+    // MTP unconditionally, which on Apple GPU generation 15+ advertised a
+    // temperature-0 byte-identity the default kernel selection does not
+    // provide (issue #1188). The gate's own retry drops `qmv_wide` when that
+    // restores exactness, so on those hosts this typically engages MTP at the
+    // exact-kernel cost rather than declining outright.
+    if let LoadedModel::Gemma4(wrapper) = model
+        && !wrapper.mtp_exactness_allows(block_size)
+    {
+        return Err(gemma4_mtp_declined(block_size));
+    }
+    if let LoadedModel::Gemma4VLM(vlm) = model
+        && !vlm.text_model.mtp_exactness_allows(block_size)
+    {
+        return Err(gemma4_mtp_declined(block_size));
+    }
+    if let LoadedModel::Gemma4Unified(unified) = model
+        && !unified.text_model.mtp_exactness_allows(block_size)
+    {
+        return Err(gemma4_mtp_declined(block_size));
+    }
+
     // Resolve the concrete target reference the drafter binds to, and reject any
     // non-MTP-capable target. Mirrors the server burst dispatch
     // (`run_mtp_burst`): bind to the same concrete Gemma 4 model the adapter

src/models/gemma4.rs

@@ -29,6 +29,9 @@ use crate::distributed::pipeline::StageExecutionOutput;
 use crate::distributed::pipeline::partial_loading::filter_weight_map;
 use crate::models::model_owned::ModelOwnedSequenceState;
 use crate::models::recurrent_snapshot::{push_i32, push_optional, restore_i32, restore_optional};
+use crate::models::speculative_exactness::{
+    BlockChainExactness, ProbeKey, compare_block_against_chain, mtp_exactness_gate,
+};
 use crate::models::switch_layers::{SwitchLinear, gather_sort};
 use mlxcel_core::cache::{
     KVCacheMode, RotatingKVCacheSnapshotState, SequenceId, SequenceStateLayout,
@@ -5116,6 +5119,133 @@ impl Gemma4Wrapper {
             .replace_internal(self.model.make_caches());
     }
 
+    /// Whether MTP may engage on this loaded checkpoint at `block_size`,
+    /// measured rather than predicted.
+    ///
+    /// The Gemma 4 arms of `mtp_capable_target` used to return `true`
+    /// unconditionally, which advertised temperature-0 byte-identity the
+    /// hardware does not always provide: on Apple GPU generation 15+ MLX
+    /// dispatches `M >= 2` affine-quantized matmuls to `qmv_wide`, whose
+    /// K-reduction order differs from the single-token `qmv`, and the
+    /// measured result is a systematic token divergence, not f16 jitter
+    /// (issue #1188). This routes through the same
+    /// [`mtp_exactness_gate`] the Qwen 3.5 family uses: probe once per
+    /// process at the configured block width, on failure retry with
+    /// `qmv_wide` disabled and keep it off when that restores exactness
+    /// (measured ~23% on this family's verify forward, #1188), and decline
+    /// to classic decode when neither arm is exact unless
+    /// `MLXCEL_MTP_ALLOW_INEXACT` is set.
+    ///
+    /// Used by: `mtp_capable_target` (server burst dispatch) and the
+    /// offline CLI MTP gate in `commands::generate`.
+    pub fn mtp_exactness_allows(&self, block_size: usize) -> bool {
+        let key = ProbeKey {
+            block_size: block_size as u32,
+            hidden_size: self.model.config.hidden_size as u32,
+            num_hidden_layers: self.model.config.num_hidden_layers as u32,
+        };
+        mtp_exactness_gate(key, || self.probe_block_chain_exactness(block_size))
+    }
+
+    /// Measure whether a `T = block_size` verify block produces
+    /// byte-identical logits to `block_size` single-token decode steps on
+    /// this checkpoint.
+    ///
+    /// Mirrors `Qwen35Model::probe_block_chain_exactness`; the comparison,
+    /// the multi-draw rationale, and the failure policy live in
+    /// [`crate::models::speculative_exactness`]. Runs on throwaway caches
+    /// built directly from the inner model, so neither the wrapper's
+    /// fallback `internal` slot nor any scheduler-owned `seq_id` slot is
+    /// touched — the constraint that kept the Gemma 4 arms on an
+    /// unconditional `true` before this existed.
+    ///
+    /// Both arms project full-width logits through the tied LM head
+    /// (`forward_with_caches_and_embeddings`), so the probe covers the
+    /// `M = K` versus `M = 1` dispatch of the head projection as well as
+    /// the decoder layers, exactly as the Qwen probe does.
+    pub fn probe_block_chain_exactness(&self, block_size: usize) -> BlockChainExactness {
+        // Mirrors `PROBE_PROMPT_LEN` / `PROBE_DRAWS` in `models::qwen3_5`.
+        const PROBE_PROMPT_LEN: usize = 8;
+        const PROBE_DRAWS: usize = 3;
+
+        if block_size < 2 {
+            return BlockChainExactness::NotRun("block width below 2 drafts nothing");
+        }
+        let vocab = self.model.config.vocab_size;
+        if vocab < 2 {
+            return BlockChainExactness::NotRun("degenerate vocabulary");
+        }
+
+        let as_input =
+            |tokens: &[i32]| mlxcel_core::from_slice_i32(tokens, &[1, tokens.len() as i32]);
+        let position_bytes = |logits: &UniquePtr<MlxArray>, index: i32| -> Vec<u8> {
+            let shape = mlxcel_core::array_shape(logits);
+            let row = mlxcel_core::slice(logits, &[0, index, 0], &[shape[0], index + 1, shape[2]]);
+            mlxcel_core::array_to_raw_bytes(&row)
+        };
+
+        for draw in 0..PROBE_DRAWS {
+            // Synthetic ids, varied per draw: dispatch depends only on shape
+            // and `M`, but whether a last-ulp kernel difference lands on a
+            // differing byte depends on the values (see the false-pass note
+            // in `models::speculative_exactness`).
+            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();
+
+            // Chain arm: prefill, then one token at a time — the shape
+            // classic decode runs, and the contract's reference.
+            let mut chain_caches = self.model.make_caches();
+            let _ = self.model.forward_with_caches_and_embeddings(
+                &as_input(&prompt),
+                None,
+                &mut chain_caches,
+                None,
+                None,
+            );
+            let mut chain_positions: Vec<Vec<u8>> = Vec::with_capacity(block_size);
+            for token in &block {
+                let out = self.model.forward_with_caches_and_embeddings(
+                    &as_input(&[*token]),
+                    None,
+                    &mut chain_caches,
+                    None,
+                    None,
+                );
+                chain_positions.push(position_bytes(&out, 0));
+            }
+
+            // Block arm: fresh caches, same prefill, the whole block at once.
+            let mut block_caches = self.model.make_caches();
+            let _ = self.model.forward_with_caches_and_embeddings(
+                &as_input(&prompt),
+                None,
+                &mut block_caches,
+                None,
+                None,
+            );
+            let out = self.model.forward_with_caches_and_embeddings(
+                &as_input(&block),
+                None,
+                &mut block_caches,
+                None,
+                None,
+            );
+            let block_positions: Vec<Vec<u8>> = (0..block_size)
+                .map(|i| position_bytes(&out, i as i32))
+                .collect();
+
+            let verdict = compare_block_against_chain(&block_positions, &chain_positions);
+            if !verdict.is_equal() {
+                return verdict;
+            }
+        }
+        BlockChainExactness::Equal
+    }
+
     pub(crate) fn input_embeddings(&self, input_ids: &MlxArray) -> UniquePtr<MlxArray> {
         self.model.text_model.embed_tokens.forward(input_ids)
     }

src/models/speculative_exactness.rs

@@ -258,8 +258,9 @@ where
             "MTP exactness probe failed under qmv_wide ({}) and passed without it. \
              Disabling qmv_wide for this process to keep the temperature-0 \
              byte-identity contract; the verify forward costs about 17 to 20 \
-             percent more. Set MLXCEL_QMV_WIDE=1 to pin the faster kernel and \
-             decline MTP instead.",
+             percent more on the Qwen 3.5 family and about 23 percent on \
+             Gemma 4 (#1188). Set MLXCEL_QMV_WIDE=1 to pin the faster kernel \
+             and decline MTP instead.",
             first.reason()
         );
         Some(true)

src/server/batch/speculative_burst.rs

@@ -682,7 +682,16 @@ fn ragged_target_sliding_window(model: &LoadedModel) -> Option<usize> {
 /// used to check a different subset of it.
 pub(crate) fn mtp_capable_target(model: &LoadedModel, block_size: usize) -> bool {
     match model {
-        LoadedModel::Gemma4(_) | LoadedModel::Gemma4VLM(_) | LoadedModel::Gemma4Unified(_) => true,
+        // The Gemma 4 arms used to return `true` unconditionally, which
+        // advertised a byte-identity the hardware does not always provide:
+        // measured on generation 15+, the default `qmv_wide` dispatch makes
+        // the verify block diverge from the chain systematically (#1188).
+        // Same gate as the Qwen arms below: probe once, buy exactness back
+        // by dropping `qmv_wide` where that suffices (~23% on this family's
+        // verify forward), decline otherwise.
+        LoadedModel::Gemma4(m) => m.mtp_exactness_allows(block_size),
+        LoadedModel::Gemma4VLM(vlm) => vlm.text_model.mtp_exactness_allows(block_size),
+        LoadedModel::Gemma4Unified(unified) => unified.text_model.mtp_exactness_allows(block_size),
         LoadedModel::Qwen35(m) | LoadedModel::Qwen35Moe(m) => m.mtp_exactness_allows(block_size),
         LoadedModel::Qwen35VLM(vlm) | LoadedModel::Qwen35MoeVLM(vlm) => {
             vlm.text_model.mtp_exactness_allows(block_size)