Round MLX_SDPA_BLOCKS up to a multiple of 32

[pierre427]

Summary

MLX_SDPA_BLOCKS (the 2-pass vector-SDPA block-count override added in #3455) is
validated only for > 0, but the kernel's second pass assumes the block count is a
multiple of 32. Any other value silently corrupts the attention output on every decode
step — no error, no clamp. This PR rounds the override up to the next multiple of 32 and
adds a regression test.

Root cause

scaled_dot_product_attention.cpp applies the override as-is:

if (int blocks_env = env::get_var("MLX_SDPA_BLOCKS", 0); blocks_env > 0) {
  blocks = blocks_env;
}

but the pass-2 reduction loops in sdpa_vector.h iterate blocks / BN (BN = 32) with
integer division:

for (int b = 0; b < blocks / BN; ++b) {
  max_score = max(max_score, maxs[simd_lid + BN * b]);
}

so for blocks % 32 != 0 the tail partials are dropped from the max/sum/output
reductions. All built-in block choices are multiples of 32 (32–1024), so only the env
override can reach the broken case.

Measured

Apple M5, macOS 26.2, decode shape B=1, 32 q-heads / 8 kv-heads, D=128, kL=8192,
fp16, max abs err vs an fp32 reference (reproduces identically on 0.32.0 (PyPI) and
current main):

MLX_SDPA_BLOCKS	max abs err
unset	2.5e-5
16	7.4e-2
31	7.4e-2
33	1.2e-2
48	4.6e-2
64	2.5e-5
100	1.3e-2

Since the env var exists as a performance-tuning lever, anyone sweeping it (say,
MLX_SDPA_BLOCKS=48) corrupts decode without noticing.

Fix

Round up rather than reject, so the knob keeps working for any positive value:

blocks = ((blocks_env + 31) / 32) * 32;

After the fix, the previously-corrupt values measure 2.1–2.5e-5 against the fp32
reference on the shape above — indistinguishable from the un-overridden path.

Alternative considered: ignore invalid values and keep the heuristic choice. Rounding
seemed friendlier for a tuning knob, but happy to switch if you prefer strict validation.

Test

test_sdpa_blocks_env_override in test_fast_sdpa.py runs the decode shape above with
MLX_SDPA_BLOCKS ∈ {16, 33, 48, 100} and asserts the output matches the un-overridden
result. All four values fail before this change and pass after (verified both ways on an
M5; the full test_fast_sdpa.py suite passes with the change).

(Noted in passing while reading this code, not addressed here: the pass-2 partials offset
arithmetic is 32-bit and would overflow once B·H·qL·blocks·D ≥ 2³¹ — reachable only at
extreme shape × blocks combinations.)

🤖 Generated with Claude Code