fast: add scaled_dot_product_attention_pad_to_fused helper

[Ogilthorp3]

Summary

Adds mlx_rs::fast::scaled_dot_product_attention_pad_to_fused — a wrapper that automatically pads head_dim up to the smallest fused-kernel size when the input falls outside MLX's supported sets, runs the fused SDPA, then slices the output back. Mirrors the ensure_fused_sdpa pattern from MLX's Python VLM examples (mlx_vlm/models/base.py).

Why

MLX dispatches optimized Metal SDPA kernels for two head_dim sets (see mlx/backend/metal/scaled_dot_product_attention.cpp:618-624):

• Prefill (Q > 1): {64, 80, 128}
• Decode (Q = 1): {64, 96, 128, 256}

Outside those sets, MLX falls back to a general-purpose attention that materializes the full (Q, K) score matrix — measurably slower at long sequence lengths. Real models that hit the fallback:

• Qwen3-VL vision tower: head_dim = 1152 / 16 = 72 → fallback at prefill.
• Several non-power-of-two LLaMA / DiT variants in the 90-110 range.

Implementation

Pure pre-existing-op composition — no new C bindings or kernels. Pads q/k/v on the last axis to the next supported size, dispatches the regular scaled_dot_product_attention, slices the result back. The padded slots contribute zero to the dot product (q·k^T = 0) and zero to the post-softmax weighted-V sum, so the result is mathematically equivalent up to fp16/fp32 rounding.

Pass-through behavior: when head_dim is already in the fused set, OR when no supported size fits (head_dim > 256 in decode, > 128 in prefill), this calls plain scaled_dot_product_attention with no pad/slice round-trip.

Validation

10 fast::tests pass under cargo test --release -p mlx-rs --lib fast::tests -- --test-threads=1, including new tests:

• test_next_fused_head_dim_decode — covers the {64, 96, 128, 256} lookup.
• test_next_fused_head_dim_prefill — covers the {64, 80, 128} lookup.
• test_pad_to_fused_matches_unpadded_at_head_dim_72 — numerically asserts the padded path produces the same output as the unpadded call (max abs diff < 1e-4) on a Qwen3-VL-shaped tensor.
• test_pad_to_fused_passthrough_when_already_fused — confirms bit-identical output when head_dim=64 (no pad triggered).

Benchmark

cargo run --release -p mlx-rs --example sdpa_pad_to_fused_bench on Apple M4 Max (sync after each call to measure end-to-end dispatch + kernel time):

Shape (B, H, L, D)	Plain SDPA	Padded SDPA	Δ
Qwen3-VL vision @ L=1024, D=72	1.31 ms	0.87 ms	1.51× faster
Qwen3-VL vision @ L=4096, D=72	13.9 ms	9.87 ms	1.41× faster
90-dim transformer @ L=2048, D=90	7.62 ms	7.98 ms	0.96× slower (90→128 pad ratio too wide)
Step-Audio-2 encoder, D=64	1.36 ms	1.36 ms	passthrough
Llama-style, D=128	7.56 ms	7.49 ms	passthrough

Pad-ratio caveat: the helper grows head_dim by pad_to / head_dim. For ratios ≤ 1.25× (e.g. 72→80) the fused kernel reliably wins. For larger ratios (e.g. 90→128 = 1.42×) the extra kernel work can outweigh the gain. Documented in the function rustdoc with a recommendation to bench specific shapes.

Discovery context

Found while wiring Qwen3-VL multimodal support in a downstream mlx-rs consumer. The vision tower's head_dim = 72 was hitting MLX's fallback path; padding to 80 lifted the entire forward back onto the fused kernel. With this helper, the Rust API reaches parity with the Python VLM idiom and the speedup is available to any caller without re-implementing the trick.

Related

Same author's previous PR #347 (fast::rms_norm: accept Option<&Array> for weight) is also in flight — the discovery here came from the same investigation. They're independent and can land in either order.

This PR addresses one part of the perf-investigation surface in #339 (flash attention wiring for ASR/FLUX encoder paths). For Step-Audio-2 (head_dim=64) the helper is a no-op (already fused). For the vision-tower-style shapes in the issue's benchmark table, the helper closes a meaningful gap without the complexity of integrating the OminiX flash-attention kernel.

[Ogilthorp3]

Hi — friendly check-in on this one (open ~7 days, MERGEABLE, no CI runs yet). Happy to address review feedback or rebase if needed. The numbers from the included bench example on M4 Max: at Qwen3-VL shapes (head_dim=72, L=1024) the padded-to-fused path is 1.51× faster than the unpadded fallback; at L=4096 it's 1.41× faster. The math equivalence test (test_pad_to_fused_matches_unpadded_at_head_dim_72) confirms identical output to bf16 tolerance, so it's a free win for any model with head_dim ∉ {64,80,96,128,256}. Anything I can do to help move it forward?