Wire flash attention kernel into ASR/FLUX encoder paths

[ymote]

Context

A custom flash attention Metal kernel with online softmax was added in mlx-rs-core/src/metal_kernels.rs (flash_attention()). It avoids materializing the full attention matrix, supports GQA and causal masking.

Benchmarking showed it's slower than MLX SDPA for Q=1 decode (TTS), but should excel on long-prefill encoder paths (ASR encoder, FLUX spatial attention) where Q >> 1 and MLX falls back to regular operations.

TODO

• Benchmark flash kernel vs MLX SDPA at Step-Audio-2 encoder dimensions (1500 tokens, 20 heads) — need model weights
• If faster, wire into step-audio2-mlx/src/encoder.rs behind a feature flag
• Test on FLUX double-block spatial attention (4096 tokens)
• Consider integrating the now-fixed MFA library as an alternative backend

Measured dimensions where gains are expected

Model	Seq tokens	Heads	SDPA ms (24L)	Manual ms (24L)
Step-Audio-2	1500	20	59ms	119ms
Qwen3-VL	1024	16	28ms	54ms

[Ogilthorp3]

Spent some time digging into this from the encoder-path angle. A few findings worth surfacing:

1. MLX has TWO fused-SDPA head_dim sets, not one. From mlx/backend/metal/scaled_dot_product_attention.cpp:618-624:

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

Anything outside those sets falls back to a general-purpose attention that materializes the full (Q, K) score matrix. This matters for the encoder path because encoders are pure-prefill (Q = L = 1500 for Step-Audio-2), so they're locked to the {64, 80, 128} set.

2. Step-Audio-2 encoder head_dim is already fused. With n_state=1280, n_head=20, head_dim = 64 — already in both fused sets. So MLX SDPA there is already on the fastest available path; the issue's table showing "SDPA 59ms" for that geometry is the fused-kernel number.

3. The 119ms "Manual" number is the real headroom. That's from step-audio2-mlx's manual implementation, not from the MLX fallback. The custom flash_attention() Metal kernel in OminiX-MLX/mlx-rs-core/src/metal_kernels.rs would need to beat MLX's fused kernel itself (59ms), not just the manual baseline (119ms), to be worth wiring in.

4. Qwen3-VL row of the table is a different story. head_dim = 1152 / 16 = 72 falls outside the fused prefill set, so MLX is hitting the fallback for that one. The "SDPA 28ms" number is the fallback path; padding head_dim 72 → 80 to land on the fused kernel beats it.

5. Submitted #348 (scaled_dot_product_attention_pad_to_fused) — a generic helper that pads head_dim up to the next supported fused size, dispatches the fused kernel, slices back. Real measurements on M4 Max:

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
Step-Audio-2 encoder (D=64)	1.36 ms	1.36 ms	passthrough (already fused)

Concrete suggestion: the original benchmark in this issue likely re-tests at the same Step-Audio-2 dims (where MLX SDPA is already optimal). For a fairer comparison, the flash kernel should also be benched at fallback-hitting shapes (head_dim outside {64, 80, 128}) where there's actually a kernel selection gap to fill. Otherwise it's competing against MLX's best-case path.

Happy to extend the bench in #348 to include the OminiX flash_attention() if helpful — just let me know whether the benchmark should live in mlx-rs or step-audio2-mlx.