fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672

Summary

On CUDA (verified NVIDIA GB10, sm_121), gemma-4-31b-it-4bit out-of-memories during a long single-pass prefill: cudaMallocAsync(...) failed: out of memory at 32768 prompt tokens, while it runs fine at 2048/4096/8192/16384. The peak scales as O(L^2), which points to the attention score matrix being materialized rather than flashed.

Evidence

• Runs OK @2048, @4096, @8192, @16384; OOMs @32768 (single-pass prefill via mlxcel-bench-decode --prompt-tokens N, and via mlxcel generate with a ~32k-token prompt).
• The memory jump matches a materialized [heads, L, L] score buffer: 32 heads * 32768^2 * 2 bytes ~= 68 GiB at 32768 (vs ~17 GiB at 16384, the exact 4x of doubling L). Weights (~17 GiB) + KV (~2 GiB, sliding-window aware) + logits + that 68 GiB pushes total over the ~115 GiB budget.
• A 31B-4bit model should need ~40-50 GiB for a 32k prefill, so this ~115+ GiB is roughly 3x too high.
• Also observed: gemma decode is abnormally slow (~0.9 tok/s at 2048 ctx on GB10), consistent with a non-flash attention path.

Root cause

Confirmed: the cuDNN flash SDPA is disqualified by its head_dim cap of 128, and gemma-4 has head_dim = 256. supports_sdpa_cudnn in _deps/mlx-src/mlx/backend/cuda/scaled_dot_product_attention.cpp:342-346 rejects the cuDNN flash path whenever the query/value head dimension exceeds 128:

if ((q.shape(-1) % 8 != 0) || (q.shape(-1) > 128) ||
    (v.shape(-1) % 8 != 0) || (v.shape(-1) > 128)) {
  return false;
}

gemma-4's head_dim = 256 > 128, so supports_sdpa_cudnn returns false and the CUDA SDPA falls back to the non-flash path that materializes the full [heads, L, L] scores.

This is compounded by gemma-4's long prefill building a dense [L, L] windowed/causal additive mask (create_causal_mask_with_window_full in src/lib/mlxcel-core/src/utils.rs, used by causal_attention in src/lib/mlxcel-core/src/lib.rs for the over-window prefill path) and handing it to the SDPA. Note the cuDNN SDPA path DOES otherwise support an array mask as a bias (options.set_bias("BIAS", mask_arr) in the same file), so flash-with-mask is in principle available for head_dim <= 128; it is specifically the head_dim = 256 gate at lines 342-346 that disqualifies gemma here.

Impact / scope

Affects gemma-4 (and gemma-3) at long context, and likely other sliding-window / large-head_dim models (the dense prefill-mask helpers are also used by Cohere2, Gemma3n, Olmo3 per code comments). It makes long-prompt prefill OOM and slows decode. It is SEPARATE from the qmm gridDim>65535 / int-overflow fix in PR #652 (issue #648).

Suggested fix direction (to investigate)

1. Route gemma / long-context prefill attention through the flash path: use fast_scaled_dot_product_attention_causal (causal flag, no materialized mask) for full-attention layers, and the cuDNN bias-mask flash path for windowed layers, instead of a dense [L, L] mask + score materialization.
2. The cuDNN flash kernel hard-caps head_dim at 128 (scaled_dot_product_attention.cpp:342-346), so it will never accept gemma's head_dim = 256 as-is. Resolve this by routing gemma to a different memory-efficient attention that supports head_dim = 256 (e.g. the fast_scaled_dot_product_attention / steel/tiled path if it accepts 256, or a chunked/tiled attention that avoids O(L^2) score materialization), or by splitting/handling the 256-dim heads so a <=128 kernel applies.
3. Verify the fix makes gemma-4-31b-it-4bit @32768 fit and improves decode throughput, and re-check the other sliding-window models.

Acceptance criteria

• gemma-4-31b-it-4bit completes a 32768-token single-pass prefill within the GB10 memory budget (no OOM).
• Peak memory for gemma long-prompt prefill scales ~O(L), not O(L^2).
• Greedy output stays token-identical to the current (pre-change) short-context behavior for a fixed seed.