fix: Diagnose and restore Gemma 4 MTP exactness on Metal #1983

Problem

After #1978, gemma-4-31b-it-qat-4bit + gemma-4-31b-it-assistant-bf16 (block_size 4) on Apple Silicon Metal (~261 GB limit) declines MTP at startup:

verify block position 0 differs from the single-token chain in 224699 of 524288 logit bytes. Disabling qmv_wide did not make it exact either.

GB10 CUDA diverges too (226171 bytes), so the M=4 verify block and the M=1 chain take different kernels somewhere other than qmv_wide. Refs #1279 (same symptom on the non-QAT 31B, M3 Ultra, LM head hypothesis only; this issue adds a localization pass and two further candidates), #1188, and docs/benchmark_results/speculative-decoding-m1ultra-2026-08-19.md (12B op-level exact, model-level divergent).

Candidate causes (unconfirmed)

• The 10 full-attention layers (of 60) use global_head_dim 512 (src/models/gemma4.rs:311, head_dim_for_layer). sdpa_supports_fast_path_impl (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4241-4274) accepts only 64/96/128/256 (vector) and 64/80/128/256 (full), so those layers fall back to unfused matmul attention: matrix-vector at M=1, matrix-matrix at M=4, different reduction order. Sliding layers (head_dim 256, gqa 2) stay on sdpa_vector for both.
• The QAT checkpoint overrides all 180 MLP gate/up/down projections to 8 bits (group 64, affine) while the default is 4-bit (confirmed in /home/inureyes/models/mlx/gemma-4-31b-it-qat-4bit/config.json); 8-bit dispatch may differ between M=1 and M=4.
• Tied LM head / embedding projection.

Phase 0: log hygiene

With MLXCEL_MTP_ALLOW_INEXACT=1 on the same host, every verify round logs this once per sliding layer (50 layers), flooding the log on long-prompt requests:

WARN mlxcel::models::gemma4: trim_mask_to_keys: mask shorter than key length, discarding caller mask mask_len=1027 key_len=1046 query_len=4

Per the comment in trim_mask_to_keys (src/models/gemma4.rs:2545-2590, #885/#891) this is the expected buffered MTP RotatingKVCache verify layout (uncompacted sliding_live_len + query_len keys vs the window-capped mask), and falling back to causal_attention with the layer window is intended. Keep the signal for unexpected cases but stop the flood: WARN once per process then DEBUG, or DEBUG when the caller is the MTP verify path. Add a test next to trim_mask_to_keys_discards_too_short_multi_token (src/models/gemma4.rs:6552) proving the fallback's effective mask for query_len=4 over the buffered chronological layout equals the causal plus sliding-window mask the chain applies, so the discard is shown exact rather than assumed.

Round statistics from that run (prompt 6449 tokens, 513 generated, block_size 4): 247 rounds, acceptance_rate 0.359, emitted_per_verify 2.08, zero/partial/full accept rounds 94/113/40, verify_forward_ms 11031, draft_ms 1460, prefill_seed_ms 6151 (a 6.15 s single-tick head-of-line stall, burst_wall_ms 6153). Low acceptance on an inexact block is consistent with the divergence below, so record acceptance before and after Phase 2.

Phase 1: localize

When probe_block_chain_exactness (src/models/gemma4.rs:4967) diverges, re-run it with per-layer capture and report the first diverging layer index, its kind (sliding/full), and the first diverging sub-op (attention output, MLP output, or final norm/LM head) in one WARN line and in the string returned by speculative_exactness::decline_reason (src/models/speculative_exactness.rs:185). Opt-in via MLXCEL_MTP_PROBE_TRACE=1 or automatic on failure only; never on the request path. Extend the existing probe, not a parallel one. Unit-test the localization function on synthetic per-layer byte vectors.

Phase 2: fix (per Phase 1's finding)

If head_dim 512 is the cause, compute those layers' verify-block attention per query row (M=1 shapes) in Attention::forward (src/models/gemma4.rs:1858), gated to the MTP verify path only. If the 8-bit MLP or LM head is the cause, apply the analogous per-row or kernel-pinning fix, or document why exactness is unreachable.

Acceptance Criteria

• Phase 0: a long-prompt MTP verify request no longer logs the trim_mask_to_keys WARN per layer per round; unexpected short masks still surface; the effective-mask equivalence test passes. Lands in the same PR as Phase 1.
• acceptance_rate and emitted_per_verify for the run above recorded before and after Phase 2.
• Phase 1 unit test passes; a diverging probe names layer, kind, and sub-op.
• macOS host logs "MTP exactness probe passed" for this pair.
• Temperature-0 MTP output byte-identical to classic decode on >=3 prompts of >=256 tokens.
• MTP vs classic decode tok/s (n>=3) recorded in the PR; classic decode output and speed unchanged.
• GB10 result recorded if the fix is backend-neutral.

Verification

cargo test --release --lib speculative_exactness gemma4
MLXCEL_MTP_PROBE_TRACE=1 mlxcel-server -m models/gemma-4-31b-it-qat-4bit --model-draft models/gemma-4-31b-it-assistant-bf16 --draft-block-size 4

Implement all three phases in one PR, run the local gate once at the end, and keep pushes minimal.