fix: Gemma 4 MTP never engages on CUDA (GB10), verify probe fails #2160

Problem / Background

Gemma 4 MTP speculative decoding never engages on CUDA. On GB10 the startup block-vs-chain exactness probe already fails on the first 8-token short draw for both the 12B and the 31B pair, so every pairing falls back to classic decode. Fail-closed works: without MLXCEL_MTP_ALLOW_INEXACT the server declines and serves classic output. The cost is that GB10 gets no MTP speedup at all for Gemma 4.

Measured during the #1986 investigation (comment): main 7c3bece5, release, --features cuda, GB10, driver 580.178.04, kernel 7.0.0-1019-nvidia, no inexact override.

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

MLX_ENABLE_TF32=0, MLXCEL_SDPA_FALLBACK_MAX_QUERIES=0 and MLX_CUDA_USE_CUDNN_SDPA=0 change nothing. PR #1987 restored 31B exactness on Metal but was never measured on CUDA. #1986 (12B long-context parity on Metal) is a separate problem and stays open.

Current Behavior

The 12B and 31B failures have different causes and need separate fixes.

12B, attention. Geometry is 16 query heads, 8 sliding KV heads at head_dim 256, 1 global KV head at head_dim 512, sliding window 1024. Neither per-row gate in Gemma4Attention::attend (src/models/gemma4.rs:2054-2104) matches it: they are hard-coded to the 31B geometry (32/4/512 full, 32/16/256 sliding with a ring cursor). So the verify block runs M=4 through causal_attention/attention_from_ptr. CUDA's supports_sdpa_vector (src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu:696-726) only accepts query_sequence_length < 4 and no array mask, so the M=4 block takes the unfused ops fallback while classic M=1 decode at head_dim 256 takes the fused sdpa_vector kernel. That is the layer 0 / pos 0 divergence. Head_dim 512 full layers are on the fallback for both sides, but M=4 and M=1 still pick different reductions inside it (the 31B layer-41 finding behind #1987), so the full-layer gate needs widening too, not just the sliding one.

12B, MLP (visible once attention is forced equal). With LARGE_D=0 the divergence moves to the layer 0 MLP, whose projections are 8-bit affine (group 64) in this checkpoint. On sm_121 both M=1 and M=4 select qmv (M * B < 8, src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:282-289), so the branch is the same, but underneath M=2..8 goes to the mlxcel multirow kernel (#725, src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu:1-17, kill switch MLXCEL_QMV_MULTIROW=0 read once per process at line 770) while M=1 uses the stock kernel. The patch comment claims per-row bit-identity with the stock kernel; this measurement suggests that claim fails at least for 8-bit. Not yet proven: a norm or GELU sub-op could also be row-count dependent. There is no CUDA analog of Metal's qmv_wide retry (retry_without_qmv_wide in src/models/speculative_exactness.rs is Metal-gated).

31B. Both per-row gates already engage on CUDA (attend has no platform guard), so the layer 6 / pos 2 divergence is a residual left after per-row attention. Candidates: the multirow qmv above (31B is uniform 4-bit), or the per-row sdpa_vector call differing from classic's sdpa_vector call on strided, sliced keys, which is the unresolved #1935 class.

Precedent. Qwen 3.5 on CUDA (PR #1944, Qwen35Model::cuda_sdpa_vector_verify_hazard, src/models/qwen3_5.rs:1630-1661) returns BlockChainExactness::NotRun for head_dim 256/288 with the fused kernel enabled. docs/benchmark_results/dflash-sdpav-options-qwen35-gb10-2026-09-21.md establishes that "fallback for verify only" cannot work (it guarantees the two sides disagree) and that LARGE_D=0 changes classic decode's output. The fix here must therefore make the verify block use the same kernels as classic, not move one side off them.

Proposed Solution

Do this in order, in one PR.

1. Attribution matrix first. On the short 8-token draw, for 12B and 31B, run the startup probe across {MLXCEL_QMV_MULTIROW unset, =0} x {MLXCEL_SDPA_VECTOR_LARGE_D unset, =0} and record the localized first divergence for each of the 8 cells. Repeat after step 2. This decides whether step 3 is needed and whether a fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode #1935-class residual exists.
2. Per-row attention for the 12B geometry on CUDA. Replace the hard-coded geometry checks in attend with a flag computed at Gemma4Attention construction from one config predicate, so each verify row is an M=1, maskless call over the physically trimmed key prefix, exactly as classic decode issues it (reuse attend_query_rows and attend_ring_rows in src/models/gemma4_attention.rs; do not write a third copy). Sliding layers must go per-row also when ring_cursor is None (short prompt, buffer not wrapped): use attend_query_rows with the layer window there. Extend TextConfig::mtp_requires_linear_singleton (src/models/gemma4.rs:267) to also return true for the 12B geometry (16/8/1, 256/512) when mlxcel_core::cuda_is_available(). That predicate already drives probe lengths (adds the long 1,056-token buffered draw), the rotating verify buffer in the probe (gemma4.rs:5260) and in serving (src/models/gemma4_mtp_target.rs:660, 808, 849, 1772), the snapshot decline (gemma4.rs:6204) and the burst dispatch (src/server/batch/speculative_burst.rs:2494). So 12B on CUDA inherits the 31B restrictions: B=1 linear verify only, no tree rounds, no buffered prompt-cache donation. Accepted; make the decline reason at gemma4_mtp_target.rs:817 ("Gemma 4 31B exact verification ...") geometry-neutral.
3. Quantized matmul. If the matrix shows MLXCEL_QMV_MULTIROW=0 removes the MLP divergence, fix the multirow kernel in qmv.cu so each row is bit-identical to the stock per-row launch, as its contract already states (same dequant to FMA order, accumulator type and final reduction for 4-bit and 8-bit). This keeps perf(cuda/quant): amortizing small-M quantized GEMM for batched decode on Blackwell #725's throughput for everyone. Rejected: disabling multirow process-wide, because it costs batched decode on every model. If the kernel cannot be made bit-identical, add a CUDA retry arm in mtp_exactness_gate mirroring retry_without_qmv_wide. That needs a new runtime bridge (set_qmv_multirow), because the switch is a static read once; adding that bridge is in scope. If multirow is not the cause, localize to the sub-op with the existing probe_trace localization and apply the per-row treatment to that op only.
4. Exit criterion if a residual remains. If after steps 2 and 3 the only remaining divergence is per-row sdpa_vector vs classic sdpa_vector (it moves with LARGE_D and nothing else), do not ship a partial gate. Add a Gemma analog of cuda_sdpa_vector_verify_hazard that returns NotRun with the exact reason, attach the matrix as a negative result in docs/benchmark_results/, and report back. That is the fallback exit, not the goal: the acceptance criteria below require the probe to pass.

Scope

In scope: src/models/gemma4.rs (attend, TextConfig predicate, probe), src/models/gemma4_attention.rs, src/models/gemma4_mtp_target.rs (reason strings), src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu if step 3 applies, tests, and a benchmark record under docs/benchmark_results/.

Out of scope: Metal 12B long-context parity (#1986). Qwen 3.5 CUDA decline (#1944) stays as is. Batched or tree MTP for 12B.

Implementation Notes

• Metal: gate the 12B widening on cuda_is_available() so Metal is byte-unchanged. fix: investigate Gemma 4 12B long-context MTP parity on Metal #1986's comment proposes the same widening as the Metal fix candidate; record that in fix: investigate Gemma 4 12B long-context MTP parity on Metal #1986, and list any change to shared 31B code paths in test: Apple Silicon verification pass for 2026-10-06 CUDA-only changes #2158 (Apple Silicon verification pass) for M5 Max confirmation.
• Edge cases: prompt shorter than the window (no ring cursor), a wrapped ring (long draw), partial accept rollback, and a declined burst after a pass must all keep classic output unchanged.
• Error handling: any probe failure keeps the current fail-closed path and the existing MTP declined warning with the localized divergence; no new panics.
• Environment: Gemma 4 runs with CUDA graphs disabled on this host (src/loading/mod.rs:530-560); do not re-enable them as part of this work.

Acceptance Criteria

• Attribution matrix (8 cells, before and after) recorded in the PR body and the benchmark record. (PR fix(gemma4): make MTP verify byte-identical to decode on CUDA #2185; docs/benchmark_results/gemma4-mtp-cuda-gb10-2026-10-07.md)
• Startup probe passes at --draft-block-size 4 without MLXCEL_MTP_ALLOW_INEXACT for both pairs on GB10, including the short draws and the long 1,056-token buffered draw; the server log shows MTP engaged, not declined.
• Greedy (temperature 0) output with MTP is byte-identical to classic for three free-generation prompts of 256 tokens on each pair. (default-config chat; needed the history-boundary prefill split, see PR fix(gemma4): make MTP verify byte-identical to decode on CUDA #2185)
• MTP vs classic decode tok/s measured with n>=3 interleaved repeats plus a null arm (classic vs classic on the same binary, opening and closing brackets), with paired deltas and acceptance length reported; a speedup inside the null arm's spread is reported as unresolved. (12B +106% to +122%, 31B +80% to +91%, null -5.1% to +6.2%)
• New unit test in src/models/gemma4_attention_tests.rs: head_dim 256, 16 query / 8 KV heads, width 4, windowed, comparing the verify path selected for the 12B geometry against per-row M=1 causal_attention byte for byte. Fails on current code on CUDA. (gemma4_12b_sliding_verify_rows_match_decode_byte_for_byte)
• New unit test that TextConfig::mtp_requires_linear_singleton is true for the 12B geometry on CUDA and false on Metal. Fails on current code. (linear_singleton_geometry_is_31b_everywhere_and_12b_on_cuda)
• If step 3 applies: a CUDA-only test that quantized_matmul on x [1, 4, K] equals four [1, 1, K] calls byte for byte, for 4-bit and 8-bit, group 64, at Gemma 12B MLP shapes (3840 to 15360 and back). Fails on current code. (step 3 did not apply: multirow changed no cell; the test was added as a pin and passes)
• New #[ignore] real-model test greedy_parity_mtp_gemma4_12b in tests/speculative_parity.rs, mirroring greedy_parity_mtp_gemma4_31b (add the 12B pair to REACHABLE_PAIRINGS); both pass on GB10. (the 12B pair and greedy_parity_mtp_gemma4_unified_12b already existed; its draft dir is fixed to resolve on Linux; both pass)
• Qwen recheck: greedy_parity_dflash_qwen35_4b (qwen3.5-4b-4bit + qwen3.5-4b-dflash) still declines on CUDA with the fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode #1935 reason, unchanged. No Qwen 3.8 MTP drafter is local; say so if unmeasured. (Qwen 3.8 unmeasured)
• Metal behavior unchanged, or every Metal-reaching change listed in test: Apple Silicon verification pass for 2026-10-06 CUDA-only changes #2158. (Metal-reaching changes listed in PR fix(gemma4): make MTP verify byte-identical to decode on CUDA #2185 for test: Apple Silicon verification pass for 2026-10-06 CUDA-only changes #2158)

Verification

Build unlocked; lock only the GPU-using run. CUDA tests must run single-threaded.

cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features cuda -- -D warnings
cargo build --release --features cuda --bins
cargo test --release --features cuda --lib --no-run
gpu-lock run --tag gemma4-mtp-cuda -- cargo test --release --features cuda --lib gemma4 -- --test-threads=1
gpu-lock run --tag gemma4-mtp-cuda -- cargo test --release --features cuda --test speculative_parity -- --ignored --test-threads=1 greedy_parity_mtp_gemma4 greedy_parity_dflash_qwen35_4b
# Startup probe per pair (pass = log shows MTP engaged and no "MTP declined" line):
gpu-lock run --tag gemma4-mtp-probe -- ./target/release/mlxcel-server --model /home/inureyes/models/mlx/gemma-4-12b-it-4bit --model-draft /home/inureyes/models/mlx/gemma-4-12b-it-assistant-4bit --draft-kind mtp --draft-block-size 4 --parallel 1

For throughput, reuse the method in docs/benchmark_results/dflash-sdpav-options-qwen35-gb10-2026-09-21.md and its driver docs/benchmark_results/data/dflash-sdpav-options-gb10-2026-09-21/harness/price_options.py (host gate, classic brackets, n=3 per arm). Record kernel, driver, MLX pin and commit in the report.