fix: investigate Gemma 4 12B long-context MTP parity on Metal

Problem / Background

While investigating #1983 on M5 Max (128 GiB, Metal, test-fast), an experimental long-context startup probe for models/gemma-4-12b-it-4bit with models/gemma-4-12b-it-assistant-4bit found a block-versus-chain mismatch after a 1,056-token prefill with speculative rotating-cache buffering. At block width 4, position 0 differs in 275,832/524,288 logit bytes; failure localization first reports layer 0 sliding-attention output. Disabling qmv_wide does not resolve it. This diagnostic result does not establish free-running error severity or its first introducing commit.

The new long buffered probe in #1983 is limited to its measured 31B geometry; 12B retains its existing short startup checks. Passing those short checks must not be interpreted as long-context parity. The 31B correction must not be copied to 12B solely because both have 512-wide global heads.

Proposed Solution

Reproduce with the 12B pair above, --draft-kind mtp --draft-block-size 4 --parallel 1, no inexact override, and a diagnostic probe that adds the 1,056-token buffered draw used during #1983. Compare physical rotating-ring ordering, mask/dispatch geometry, and prefix state against classic M=1 decode; quantify the mismatch before selecting a correction. Keep the default startup policy change separate until evidence supports it.

Acceptance Criteria

Reproduce on the current baseline and record provenance, actual buffered cache layout, and first divergent sub-operation.
Report teacher-forced top-1 and decided-position disagreements at realistic context and actual verify widths, plus three free-generation comparisons.
Correct verified causes and extend the 12B startup guard with measured startup cost; retain fail-closed behavior without an inexact override.
Recheck 31B and Qwen MTP and measure classic/MTP performance; document untested backends.

docs: add technical report for PR #1987

fix: restore Gemma 4 31B MTP exactness (#1987)

CUDA (GB10) result: does not isolate the Metal defect

Provenance: main 7c3bece5, release profile, --features cuda, NVIDIA GB10, driver 580.178.04, kernel 7.0.0-1019-nvidia, MLX_ENABLE_TF32 at its default unless noted. Checkpoints gemma-4-12b-it-4bit + gemma-4-12b-it-assistant-4bit, and gemma-4-31b-it-4bit + gemma-4-31b-it-assistant-bf16. No inexact override in any run. The CUDA qmv_wide retry does not exist (Metal-gated), so no decline below includes a narrow retry.

Command (startup probe, killed after /health):

mlxcel-server --model TARGET --model-draft DRAFT --draft-kind mtp --draft-block-size 4 --parallel 1

The long draw needs the diagnostic patch at the end of this comment: MLXCEL_DIAG_GEMMA4_LONG_PROBE=1 adds the 1,056-token draw and enables the MTP rotating buffer on the block arm (as #1987 does for 31B); MLXCEL_DIAG_GEMMA4_PROBE_LENGTHS=a,b,... overrides the draw list; each draw prints a [DIAG1986] line.

Pair, draws	Failing draw	Bytes differing (of 524,288)	First divergence
12B, unpatched default (8,8,8)	draw 0, prompt 8	299,412	layer 0 sliding, attention output, pos 0
12B, 1056 only	draw 0, prompt 1056	241,285	layer 0 sliding, attention output, pos 0
31B, unpatched default (8,8,8,1056)	draw 0, prompt 8	243,000	layer 6 sliding, attention output, pos 2
12B, 8 only, `MLXCEL_SDPA_VECTOR_LARGE_D=0`	draw 0	269,030	layer 0 sliding, MLP output, pos 0
12B, 1056 only, `MLXCEL_SDPA_VECTOR_LARGE_D=0`	draw 0	251,689	layer 0 sliding, attention output, pos 1

MLX_ENABLE_TF32=0, MLXCEL_SDPA_FALLBACK_MAX_QUERIES=0 and MLX_CUDA_USE_CUDNN_SDPA=0 leave the 12B rows unchanged (cuDNN rejects head_dim 256 anyway).

Reading: on CUDA the short 8-token draw already fails for both 12B and 31B, so the Metal symptom (short passes, long buffered fails) is not reproduced in isolation. The CUDA failure is CUDA-wide T>1 versus T=1 dispatch: M=1 decode at head_dim 256 takes sdpa_vector, the M=4 verify block takes the unfused ops fallback. Forcing both onto the fallback moves the short draw's divergence to the layer-0 MLP (quantized matmul M=4 versus M=1), and moves the long draw's to position 1, which is consistent with, but does not prove, a ring-order component underneath. Fail-closed works: every pairing declines to classic decode at startup. Startup to /health: 13.2 s (short) and 11.1 s (1056 only), both including the failure-localization rerun, so neither is a clean long-draw cost.

Candidate Metal cause (code reading, unverified): 12B is 16 query heads, 8 sliding KV heads, 1 global KV head. Neither geometry gate in Gemma4Attention::attend (32/4/512 full, 32/16/256 sliding) engages, so with the buffer on the 12B block arm runs batched attention with a mask over chronological buffered keys while classic decode reduces maskless over the physical ring. That is the mechanism #1987 measured on 31B. Hypothesis only.

What a Mac run must collect

Apply the patch, MLXCEL_QMV_WIDE=0, run the command with MLXCEL_DIAG_GEMMA4_LONG_PROBE=1 and MLXCEL_DIAG_GEMMA4_PROBE_LENGTHS=8,8,8,1056, then =1056 alone. Record the [DIAG1986] lines and the localized first divergence.
Temporarily widen the sliding attend_ring_rows gate to 16/8/256 and rerun the 1056 draw. If position 0 passes and the divergence moves to a full-attention layer, also widen the full-attention gate to 16/1/512 and rerun.
Only if all draws pass: the MLXCEL_TRACE_GEMMA4 1,536-prefix trace at width 4 with --decided 2 (docs/benchmarks.md), three 256-token free-generation comparisons against classic, startup time, and a recheck of 31B and Qwen MTP.

The CUDA-wide Gemma 4 MTP decline is a separate problem and is not part of this issue.

--- a/src/models/gemma4.rs
+++ b/src/models/gemma4.rs
@@ -199,6 +199,10 @@ fn default_partial_rotary_factor() -> f32 {
     1.0
 }
 
+fn diag_long_probe() -> bool {
+    std::env::var("MLXCEL_DIAG_GEMMA4_LONG_PROBE").is_ok_and(|v| v == "1")
+}
+
 fn default_rope_type() -> String {
@@ -275,8 +279,11 @@ impl TextConfig {
     pub(crate) fn mtp_probe_prompt_lengths(&self) -> Vec<usize> {
+        if let Ok(list) = std::env::var("MLXCEL_DIAG_GEMMA4_PROBE_LENGTHS") {
+            return list.split(',').filter_map(|v| v.trim().parse().ok()).collect();
+        }
         let mut lengths = vec![8; 3];
-        if self.mtp_requires_linear_singleton() {
+        if self.mtp_requires_linear_singleton() || diag_long_probe() {
             lengths.push(self.sliding_window.saturating_add(32).min(1536));
@@ -5259,7 +5266,7 @@ impl Gemma4Wrapper {
             for cache in block_caches
                 .iter_mut()
-                .filter(|_| self.mtp_requires_linear_singleton())
+                .filter(|_| self.mtp_requires_linear_singleton() || diag_long_probe())
             {
@@ -5305,6 +5312,7 @@ impl Gemma4Wrapper {
             let verdict = compare_block_against_chain(&block_positions, &chain_positions);
+            eprintln!("[DIAG1986] draw={draw} prompt_len={prompt_len} trace={trace} verdict={}", verdict.reason());
             if !verdict.is_equal() {

test(gemma4): skip non-row-wise pairs in the batched-vs-classic test

Included in the 2026-10-12 triage and routed to macOS together with #1599 and #2158. The existing CUDA investigation explicitly does not isolate the Metal defect, so another GB10 run cannot close this issue. On current Metal main, reproduce the 12B long buffered probe using the diagnostic procedure in the 2026-10-06 comment, localize the first divergence, and widen attention geometry only if measured evidence supports it. Completion still requires teacher-forced and free-generation comparisons, startup cost, and 31B/Qwen regression checks. No speculative geometry change or completion claim is made by the GB10 batch.

update: advance MLX pin and reconcile backend ownership (#2304)

fix: scope Gemma 4 12B parity policy to supported backends

docs: record Gemma 4 parity release costs and reproduction

docs: verify final Gemma parity release and default startup

Implemented in #2315; final review and CI remain before merge.

The M1 Ultra reproduction isolates sliding-ring reduction order, maskless full-attention dispatch and serving prefill chunking. Corrected verification matches 256 teacher-forced positions and three long-prompt completions; final-source release repeats all three classic/MTP hashes with positive draft counters. A separate default-environment startup passes without an inexact override. Existing 31B/CUDA policy is preserved, and only active-Metal 12B is newly routed.

The benchmark record includes artifact provenance, exact requests, regression results, repeated release dispersion and startup cost. Forced-width MTP is slower on M1: paired ratios 0.837x/0.903x/0.815x; scheduler-setup-to-probe-pass increases by 5.795s. Defaults are not broadened. M5, CUDA, ROCm and CPU runtime validation remain unmeasured; this is not an original-M5 reproduction or CUDA-only resolution.

docs: add Gemma 4 MTP parity report

Navigation Menu

Repository navigation

fix: investigate Gemma 4 12B long-context MTP parity on Metal #1986

Problem / Background

Proposed Solution

Acceptance Criteria

CUDA (GB10) result: does not isolate the Metal defect

What a Mac run must collect

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

fix: investigate Gemma 4 12B long-context MTP parity on Metal #1986

Description

Problem / Background

Proposed Solution

Acceptance Criteria

Activity

inureyes commented on Oct 6, 2026

CUDA (GB10) result: does not isolate the Metal defect

What a Mac run must collect

inureyes commented on Oct 12, 2026

inureyes commented on Oct 14, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions