perf(cuda): decode for head_dim > 128 models runs the unfused SDPA fallback; extend the sdpa_vector kernels to cover them #675

Summary

On the CUDA backend, every decode step of a head_dim > 128 model runs MLX's unfused materializing SDPA fallback because neither fused path accepts the configuration, making decode several times slower than comparable-size llama-family models. Issue #672 fixed the prefill memory blowup for these configurations; this issue is the decode throughput follow-up.

Root cause

Fetched MLX source, mlx/backend/cuda/: supports_sdpa_cudnn (scaled_dot_product_attention.cpp, the head_dim gate at lines ~342-346) requires head_dim % 8 == 0 and head_dim <= 128 for both Q and V, and supports_sdpa_vector (scaled_dot_product_attention.cu, ~line 661) only accepts head_dim in {64, 96, 128} with q_len < 4 and no array mask. A head_dim 256 decode step therefore falls back to the explicit matmul + mask + softmax + matmul composite, with per-token kernel-launch overhead and no fused softmax.

Affected models

Local checkpoint scan, 43 of 158 flash-ineligible: the whole gemma family (gemma, gemma2, gemma3, gemma3n, gemma-4 dense/MoE/unified, diffusiongemma, all head_dim 256; paligemma2 is 288), qwen3.5 and qwen3.6 families (head_dim 256, full attention), baichuan-m1 (256), and the MLA families (deepseek-v2/v3 and kimi-vl with qk head_dim 192 / v 128, glm-5 and glm4-flash with 256, deepseek-v4-flash with 512).

Measured impact

GB10, release, benchmarks/cuda_gb10_gemma_longprompt_672_2026-07-06.csv and the 2026-07-04 sweep: gemma-4-31b-it-4bit decodes at 5.75 to 7.3 tok/s and gemma-4-12b-it-4bit at ~14 tok/s, versus llama-3.1-8b-4bit at 43 to 54 tok/s on the same hardware; the gap is far larger than the parameter-count ratio suggests.

Fix direction

Extend the CUDA sdpa_vector 1-pass and 2-pass kernels to head_dim 256 (and evaluate 192 for the MLA qk dimension and 288 for paligemma2), then widen the supports_sdpa_vector gate accordingly. The mlxcel-owned MLX overlay mechanism is already proven for exactly this kind of surgical CUDA fix by PR #652 (src/lib/mlx-cpp/patches/mlx/backend/cuda/, applied by the build over the fetched MLX source), so the kernels can be patched in the overlay first and proposed upstream afterward. Note the 2-pass variant is selected for k_len > 1024, which is the common long-context decode case for full-attention layers; sliding-window layers decode against at most window keys and will mostly take the 1-pass variant.

Validation

• Greedy token identity against the current fallback on gemma-4 and deepseek-v2-lite at short and long context.
• Decode tok/s before and after on gemma-4-12b/31b at 2048/8192/32768 context via mlxcel-bench-decode (target at least 2x on gemma-4).
• Confirm no regression for head_dim <= 128 models.
• Confirm prefill paths are untouched (q_len < 4 gate).

Acceptance criteria

• head_dim 256 decode no longer routes through the materializing fallback on CUDA
• Measured decode speedup reported in a benchmarks CSV
• Token-identity validation documented
• Overlay patch applies cleanly against the pinned MLX commit

Related

• Follow-up of fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672 (prefill memory blowup fix for the same head_dim > 128 configurations)
• Sibling of the fix(cuda/quantized): chunk long-prompt qmm launches to avoid gridDim>65535 and int overflow #652 overlay mechanism (chunked long-prompt qmm launches via the same MLX CUDA patch path)