From ef562baebc9305fd2ffc8c6ad16872c9be5f450b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 20 Aug 2026 08:33:45 +0900 Subject: [PATCH] fix(speculative): probe Gemma 4 MTP exactness instead of assuming it 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 --- docs/benchmarks.md | 29 ++++-- src/commands/generate.rs | 36 +++++++ src/models/gemma4.rs | 130 ++++++++++++++++++++++++++ src/models/speculative_exactness.rs | 5 +- src/server/batch/speculative_burst.rs | 11 ++- 5 files changed, 198 insertions(+), 13 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 2c7c800f5..191851b53 100644 --- a/docs/benchmarks.md +++ b/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: diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 0ce2ab10c..80189d2af 100644 --- a/src/commands/generate.rs +++ b/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 diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 2026ff1e1..0f07d6ade 100644 --- a/src/models/gemma4.rs +++ b/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, index: i32| -> Vec { + 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 = (0..PROBE_PROMPT_LEN).map(|i| wrap(i, 7, 1)).collect(); + let block: Vec = (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::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> = (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 { self.model.text_model.embed_tokens.forward(input_ids) } diff --git a/src/models/speculative_exactness.rs b/src/models/speculative_exactness.rs index e8aea4ff4..7605bae88 100644 --- a/src/models/speculative_exactness.rs +++ b/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) diff --git a/src/server/batch/speculative_burst.rs b/src/server/batch/speculative_burst.rs index 40e5652b9..a004310c3 100644 --- a/src/server/batch/speculative_burst.rs +++ b/src/server/batch/speculative_burst.rs @@ -682,7 +682,16 @@ fn ragged_target_sliding_window(model: &LoadedModel) -> Option { /// 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)