[inureyes]

Gemma 4 MTP now engages on CUDA. On GB10 the startup block-vs-chain probe passes for both the 12B and the 31B pair at --draft-block-size 4, including the 1,056-token buffered draw, and default-config greedy chat output is byte-identical to classic decode. Before this PR both pairs declined at the first 8-token draw.

Attribution

Full matrix, before and after, in docs/benchmark_results/gemma4-mtp-cuda-gb10-2026-10-07.md. MLXCEL_QMV_MULTIROW changed nothing in any of the cells, so the multirow qmv kernel is not a cause and step 3 (kernel fix or retry arm) was not needed. Three causes were found:

• 12B attention gate. The per-row verify path in attend only matched the 31B geometry, so the 12B M=4 block left the single-query sdpa_vector kernel decode uses at head_dim 256. TextConfig::mtp_requires_linear_singleton now also covers the 12B geometry (16/8/1 heads, 256/512) on CUDA, and attend reads a flag computed from it at construction (attend_verify_rows, reusing attend_query_rows/attend_ring_rows). The 12B pair on CUDA inherits the 31B restrictions: B=1 linear verify, no tree rounds, no buffered snapshot donation. The decline reason is now geometry-neutral.
• 8-bit compiled GeGLU. Decode ran the compiled MLP graph, a multi-token block the eager activation, which rounds differently. On CUDA, inputs under 8 rows (the qmv window) now take the compiled graph.
• RoPE kernel choice. MLX CUDA rotates B=1, L=1 with rope_single and wider inputs with the general kernel; they differ in the last float bit. Row-wise verify on CUDA rotates Q and K one row at a time through decode's exact chain.

The issue's step 4 exit (NotRun decline) was not needed: no sdpa_vector-only residual remained.

Chat history boundary

After those fixes, default-config chat still diverged on one of three 256-token prompts per pair. Classic serving splits a chat prefill at the history boundary (#1143); the MTP burst bypassed that and prefilled in one piece. With --no-cache-prompt, and via /v1/completions, all prompts matched. The scheduler now passes the boundary (history_boundary_split) to the B=1 MTP slice and burst, and the row-wise prefill splits there (mtp_prefill_ranges). After that, 3 of 3 match on both pairs.

Verification (GB10, release, --features cuda)

• Startup probe, default draws (8, 8, 8, 1056), both pairs, all four MULTIROW x LARGE_D cells: pass.
• Chat, temperature 0, 256 tokens, 3 prompts per pair, default config: byte-identical, MTP engaged (probe passed, bursts ran).
• speculative_parity --ignored (full suite, branch rebased on main 9c0ae2e9): 7 of 8 pass. The failure is b1_batched_baseline_probe, which also fails on origin/main 9c0ae2e9 without this PR (row 1 at token 1, reproduced twice); see below. Passing: greedy_parity_mtp_gemma4_31b, greedy_parity_mtp_gemma4_unified_12b (the issue's 12B parity test; the pairing existed, its draft dir is fixed to the lowercase name so it resolves on Linux), both batched Gemma 4 tests, and greedy_parity_dflash_qwen35_4b pass. Qwen 3.5 DFlash still declines with the fix(speculative): Qwen 3.5 DFlash verify diverges from classic decode #1935 reason, unchanged.
• New tests: gemma4_12b_sliding_verify_rows_match_decode_byte_for_byte, linear_singleton_geometry_is_31b_everywhere_and_12b_on_cuda, qmv_multirow_matches_per_row_qmv_bitwise_gemma4_12b_mlp_shapes, compiled_qgelu_mlp_small_block_rows_match_single_token_bitwise, mtp_prefill_ranges_mirror_classic_history_boundary_split. The gemma4 lib suite (250), scheduler burst/slice/prompt-cache selectors (127) and the core qmv/qgelu tests pass; clippy -D warnings on both crates and cargo fmt --check are clean.
• Throughput (3 interleaved rounds, classic-vs-classic null arm, prompt cache off, 200 tokens): 12B MTP +106% to +122% over classic (13.2 to 28 tok/s), 31B +80% to +91% (8.3 to 15.8 tok/s); null spread -5.1% to +6.2%; acceptance 0.63 and 0.64; identical output bytes across every arm of each pair. The CPU-quiet host gate was skipped because the shared CI runner stayed busy; details in the record.

Pre-existing: b1_batched_baseline_probe

This 31B diagnostic compares the batched adapter run at B=1 against the B=1 linear adapter. On main row 1 differs; on this branch rows 1, 2 and 3 differ (tokens 1, 9, 17). The B=1 linear adapter is now decode-exact (probe, chat parity, the parity tests); the batched adapter never takes the row-wise verify path and serving never uses it for these geometries (B>1 is declined for linear-singleton targets). So the wider mismatch is the reference moving onto decode's kernels. Making the batched adapter row-wise is outside #2160; noted in the benchmark record.

Shared code that reaches Metal (for #2158)

• attend gate refactor: same 31B behavior, plus a new per-row fallback for a sliding layer with no ring cursor while keys fit the window (unreachable in serving and the probe, which always enable the buffer).
• The history-boundary prefill split in the row-wise prefill, which Metal 31B serving also takes.
• mtp_requires_linear_singleton now consults cuda_is_available(); the Metal answer is unchanged.

The 12B predicate arm, the RoPE change and the compiled-GeGLU gate are CUDA-only.

Not verified on this host

Metal behavior (no Apple hardware here). No local Qwen 3.8 MTP drafter, so that pairing is unmeasured. The first bonus token still comes from an M=prompt LM-head projection in MTP prefill against classic's last-position projection; it matched in every run but is not covered by construction.

Closes #2160