perf(cuda): extend sdpa_vector decode kernels to head_dim 256/288 via MLX overlay

[inureyes]

Summary

Extend the CUDA sdpa_vector decode kernels and the supports_sdpa_vector gate to head_dim 256 and 288 via the mlxcel MLX source overlay, so head_dim > 128 models (gemma family 256, qwen3.5/3.6 256, baichuan-m1 256, paligemma2 288) take the fused vector decode path instead of MLX's unfused materializing SDPA fallback.

Root cause and single-file overlay approach

On CUDA, decode (q_len < 4) for head_dim > 128 always took MLX's materializing SDPA fallback: supports_sdpa_cudnn requires head_dim <= 128 and supports_sdpa_vector only accepted head_dim in {64, 96, 128}, so 256/288 fell through to the explicit matmul + mask + softmax + matmul composite with per-token launch overhead and no fused softmax.

This is a single, full-file overlay of mlx/backend/cuda/scaled_dot_product_attention.cu, copied over the fetched MLX source at CMake configure time (the same overlay mechanism as PR #652). The .cpp is deliberately untouched: ScaledDotProductAttention::use_fallback returns !supports_sdpa_cudnn && !supports_sdpa_vector and cudnn rejects head_dim > 128, so widening supports_sdpa_vector alone is sufficient to route 256/288 decode to the vector kernels. The three kernels (kernel_sdpav_1pass, kernel_sdpav_2pass_1, kernel_sdpav_2pass_2) already parameterize on template int D with fixed-size, D-independent shared tiles and a 32x32 transpose aggregation, so D=256 (v_per_thread=8) and D=288 (v_per_thread=9) are handled by the existing kernel bodies; only the dispatch and the gate needed widening.

Changes: (A) dispatch_headdim gains case 256 and case 288 plus a default: that throws std::runtime_error naming the head_dim, since upstream had no default and a gate/dispatch desync would silently launch no kernel and leave the output buffer as uninitialized garbage; this turns that failure into a loud error. (B) supports_sdpa_vector accepts head_dim 256/288 behind a static env kill switch while keeping the q_len < 4, !has_arr_mask, and !output_logsumexp gates unchanged. (C) __launch_bounds__ is added to the three kernels.

Launch-bounds rationale

The D=256/288 instantiations hold roughly 3 * v_per_thread floats in per-thread registers, so without launch bounds nvcc may compile to more than 64 regs/thread and a 1024-thread launch would fail at runtime with "too many resources requested for launch". __launch_bounds__(1024) on kernel_sdpav_1pass and kernel_sdpav_2pass_2 and __launch_bounds__(256) on kernel_sdpav_2pass_1 (their exact launch sizes) cap registers, spilling if needed, and are a no-op for D <= 128 which already fit.

Kill switch

static bool large_d_enabled = env::get_var("MLXCEL_SDPA_VECTOR_LARGE_D", 1) (default on) gates only the 256/288 additions, so setting MLXCEL_SDPA_VECTOR_LARGE_D=0 restores the exact prior materializing fallback with no rebuild; that is how the A/B below was collected and is also the rollback lever.

Correctness

Greedy decode on gemma-4-12b-it-4bit with the chat template produces byte-identical output between the vector path (default) and the fallback (MLXCEL_SDPA_VECTOR_LARGE_D=0). The vector path is provably engaged: on degenerate pad-dominated prompts the two paths diverge, which is only possible if different code runs, and that divergence is FP-ordering noise (fused online-softmax with exp2/log2e vs the materializing plain softmax) on near-tie logits, present in both directions and not a correctness defect.

Measured decode throughput (GB10, release, same bench harness and env toggle)

model	context	vector tok/s	fallback tok/s	ratio
qwen3.5-2b-4bit	8192	129.06	114.89	1.123x
qwen3.5-9b-4bit	8192	38.55	36.92	1.044x
gemma-4-12b-it-4bit	512	14.62	14.72	0.993x
gemma-4-12b-it-4bit	2048	14.33	14.32	1.001x
gemma-4-12b-it-4bit	8192	14.28	14.37	0.994x
gemma-4-12b-it-4bit	32768	13.78	13.68	1.007x

The fused path is a real +4 to +12% win on full-attention head_dim-256 families (qwen3.5/3.6, baichuan-m1), growing with attention share; full rows (prefill, peak memory, notes) are in benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv.

gemma-4 is not attention-bound (follow-up #680)

gemma-4-12b stays flat across 512/2048/8192/32768 because only 8 of its 48 layers are full-attention (a 5-sliding + 1-full pattern) and its decode is dominated by roughly 44 ms/token of model-specific fixed overhead rather than attention; that overhead, and the original 2x decode target for gemma-4, are re-attributed to follow-up issue #680, and this PR does not attempt it.

Deferred scope

Head_dim 192 and 512 are out of scope: no shipped model uses symmetric qk=v=192, and the MLA families (deepseek-v2/v3 qk 192 / v 128 and deepseek-v4-flash 512) need asymmetric D_qk != D_v kernel signatures, which is a separate change.

Overlay maintenance

The overlay is a full-file copy of upstream scaled_dot_product_attention.cu synced to MLX commit e9463bb with only the dispatch, gate, and launch-bounds deltas applied, so it must be re-synced (three-way reconciled) on any MLX pin bump; docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md now lists it in the overlay table.

Closes #675

[inureyes]

PR Finalization Complete

Summary

• Tests: not added. This PR is a single-file MLX CUDA source overlay (scaled_dot_product_attention.cu) compiled into MLX at CMake configure time; there is no Rust-reachable unit-test seam for it, CI does not build CUDA, and GPU validation (byte-identical greedy decode vs the fallback, plus the A/B throughput sweep) is already recorded in the PR body and in benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv. Adding a Rust test would not exercise the changed code path.
• Documentation: verified the overlay table row in docs/benchmark_results/mlx-pin-upgrade-2026-07-03.md renders as a well-formed table (column count matches the header) and already names the MLXCEL_SDPA_VECTOR_LARGE_D kill switch. Found docs/environment-variables.md as the existing env-var reference page and added a row for MLXCEL_SDPA_VECTOR_LARGE_D next to the other CUDA attention-path diagnostic toggles, describing the default-on behavior, what it gates, and the rollback semantics. No new pages or translations were added, per the single-source English docs convention.
• Lint/Format: verified benchmarks/cuda_gb10_sdpav_675_2026-07-06.csv has no trailing whitespace and ends with a newline, and confirmed the .cu overlay has no trailing whitespace anywhere in the file (including the added dispatch/gate/launch-bounds lines). No Rust files are touched by this PR, so cargo fmt/clippy are not applicable.

All checks passing. Ready for merge.