fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models

[inureyes]

Summary

On CUDA, gemma-4/gemma-3 (head_dim 256) disqualify both fused SDPA kernels (cuDNN caps head_dim at 128; sdpa_vector covers 64/96/128 with q_len < 4), so a long single-pass prefill fell to MLX's materializing fallback: O(heads*L^2) attention scores (~68 GiB at 32k for the 31b), full-sequence 262k-vocab logits plus a same-size final_logit_softcapping copy, and two dense [L, L] f32 masks together pushed GB10 past its ~121 GB budget and OOMed. A cuDNN head_dim cap bump alone was tested and is insufficient; the fix has to stop materializing.

Changes

• Query-chunked SDPA gate (layers::attention): when a CUDA config cannot reach a fused SDPA (head_dim > 128, softcap composite, or non-f16/bf16 dtype) and the score matrix would exceed MLXCEL_ATTENTION_CHUNK_BUDGET_MB (default 1024, 0 disables), the query axis is split into near-equal chunks. Masks are row-sliced per chunk and cast to the score dtype, because MLX's fallback adds the mask without casting and an f32 mask silently promotes the masked scores and softmax to f32. The maskless causal path in causal_attention routes through the same chunking with per-chunk causal masks and K/V sliced to the causal bound.
• Last-position LM head (LanguageModel::forward_last_logits): the single-sequence prefill in generate_streaming/generate_with_stats now requests only the sampled position's logits. The default implementation is the previous forward-plus-slice; the gemma4 override projects a single hidden row through the 262k-vocab LM head for prefills longer than 4096 tokens, dropping ~17 GiB of logits plus an equal softcap copy at 32k.
• Flag-path prefill masks for gemma4: text-only prefills longer than 4096 tokens leave the mask pair None, routing full-attention layers through the causal flag path and sliding layers through the over-window prefill path instead of retaining two dense [l, l+offset] f32 masks (~8 GiB at 32k) across the whole forward. Verify, padding, per-row valid-end, and vision-overlay shapes keep the dense masks byte-identically.
• Lean mask builders: create_causal_mask / create_causal_mask_with_window_full build from broadcast index comparisons instead of the six-buffer ones/tril/triu/multiply/greater/where chain, cutting mask-construction transients from ~24 GiB to ~7 GiB at 32k.
• Opt-in cache-level chunked prefill (MLXCEL_PREFILL_CHUNK, default 0 = single-pass): feeds the prompt through forward_last_logits in chunks, evaluating and releasing pool memory per chunk so rotating sliding-window caches trim between chunks and the cudaMallocAsync pool cannot accumulate per-shape high-water marks (measured 84 GB system peak for a ~30 GB live set without the release). Mirrors the server's prefill_chunk_size and upstream mlx-lm/mlx-vlm defaults; flipping the CLI default is feat: default the single-sequence CLI/bench prefill to chunked (2048) with a per-model opt-out policy #674.
• mlxcel-bench-decode prints the MLX allocator peak after each run; benchmarks CSVs for the post-fix(cuda/quantized): chunk long-prompt qmm launches to avoid gridDim>65535 and int overflow #652 sweep and this validation are included.

Validation (GB10 CUDA, benchmarks/cuda_gb10_gemma_longprompt_672_2026-07-06.csv)

configuration	before	after
gemma-4-12b-it-4bit @32768, single-pass	OOM class (scores alone ~34 GiB)	completes, 56.8 GB MLX peak
gemma-4-12b-it-4bit @32768, MLXCEL_PREFILL_CHUNK=2048	n/a	16.6 GB MLX peak / 23 GB system, prefill 706 tok/s
gemma-4-31b-it-4bit @16384, single-pass	~74 GB system peak	54.3 GB MLX peak / 80 GB system
gemma-4-31b-it-4bit @32768, MLXCEL_PREFILL_CHUNK=2048	OOM at ~101 GB (issue repro)	29.9 GB MLX peak / 38 GB system

Greedy output: forced-chunk A/B on gemma-4-31b is token-identical to the unchunked path on a natural prompt; short-context behavior is byte-identical by construction (the gates never fire below the budget/threshold). Decode is unchanged (13.99 vs 14.03 tok/s at 8192 over 64 tokens). 13 new unit tests cover chunked-vs-unchunked equivalence (mask, softcap GQA, broadcast mask, causal offset), chunk-length math, mask row slicing, forward_last_logits, and chunked-prefill exactly-once feeding; cargo clippy -D warnings clean on default and cuda features; the mlxcel-core suite passes with only the pre-existing environment-bound failures (verified identical on a clean tree).

Scope beyond gemma

The chunk gate automatically covers every flash-ineligible CUDA config that routes through layers::attention/causal_attention: qwen3.5/3.6 (head_dim 256), baichuan-m1 (256), paligemma2 (288 + softcap), and the MLA families (deepseek-v2/v3 qk 192, glm-5 256, deepseek-v4 512). Decode throughput for these families remains fallback-bound and is tracked as #675; defaulting the CLI to chunked prefill is #674.

Closes #672