Add MLX_SDPA_BLOCKS env var for 2-pass vector kernel block-count override

[adurham]

Summary

Adds an MLX_SDPA_BLOCKS env var that overrides the heuristic-chosen blocks (partial-tile count) for the 2-pass SDPA vector kernel. Unset or non-positive: heuristic unchanged.

The 2-pass kernel currently picks blocks from a device-class + sequence-length matrix in scaled_dot_product_attention.cpp. The defaults work well for most upstream-tested combinations but leave headroom on (device, head-count, sequence-length) tuples the heuristic doesn't anticipate. Letting operators sweep at runtime without rebuilding MLX is the value.

Empirical example

On a 2-rank M4-Ultra cluster running a 256-expert long-context MoE at K≈50k:

• Heuristic picks blocks=1024
• MLX_SDPA_BLOCKS=88 is +6.5% decode tps with a sharp cliff at 92+
• Matches the ~352-concurrent-simdgroup capacity (4 kv_heads × 88 ≈ 1.1 dispatch rounds)

Different workloads will sit at different optima. We've been carrying this in a fork for a few weeks specifically to support that kind of tuning sweep.

Files

• mlx/backend/metal/scaled_dot_product_attention.cpp — sdpa_2pass_blocks_override() helper + one call site after the existing heuristic block
• python/tests/test_fast_sdpa.py — test_sdpa_blocks_env_override: sweeps {16, 64, 256}, asserts numerics match heuristic-default output

Notes for reviewers

• Currently only wired into the bf16 2-pass path (the only one in upstream). The fork additionally applies the same override to a quantized 2-pass kernel that doesn't exist upstream — easy to extend later if you ever land that.
• I considered a build-time #define but a runtime env var is the only thing that supports sweeping without rebuilding. Happy to switch to mx.metal.set_sdpa_blocks(int) runtime API if you'd rather not key off env vars.
• Default (unset) preserves the existing heuristic exactly.

[zcbenz]

The test is not really testing anything, I'm OK with no test for the env.

[adurham]

Thanks for the review! Addressed in 4b87e47:

• Renamed helper away from override (sdpa_2pass_blocks_override → sdpa_2pass_blocks_from_env, and the local int override at the call site → int blocks_env).
• Switched to env::get_var("MLX_SDPA_BLOCKS", 0) instead of the manual std::getenv + std::atoi; dropped the now-unused <cstdlib> include.
• Dropped test_sdpa_blocks_env_override.

Left as a separate commit for review; happy to squash before merge.

[adurham]

Done in a7a77ab — helper removed, env::get_var inlined at the call site as suggested.

[zcbenz]

Looks good to me!