Add force_fused option to scaled_dot_product_attention

[DT-HJ]

Proposed changes

Add a force_fused option to mx.fast.scaled_dot_product_attention that
ignores the automatic dispatch heuristics, uses a fused kernel, and raises
with the unsupported constraint when no fused kernel is available. This lets
inference runtimes — the only party that knows their resident weights, KV
cache size, and memory budget — decide when the fused path's bounded-memory
behavior is worth its throughput cost.

Also restores the fused full-attention support for head_dim=192/256 that was
proposed in #3293/#3660. Those kernels are now reachable only through
force_fused; the default dispatch does not route to them.

Follow-up hardening incorporates the full-dispatch threadgroup safety check
and capability-diagnostic concerns raised by @apocryphx in #4186, while
retaining this PR's 192/256 kernels and C++ positional-call compatibility.

Background

#3658 tracks models with genuine head_dim=192/256 full-attention layers. At
long context, the unfused path materializes a score
transient of O(n_heads × qL × kL) per full-attention layer. The fused
full-attention kernel tiles K/V so its transient stays bounded regardless of
kL. However, benchmark measurements on a genuine head_dim=256 workload
(the Qwen 3.6 35B checkpoint measured in #3658) show the fused path trades
prompt throughput for that memory floor:

context	main prompt_tps	fused prompt_tps	main peak	fused peak
32K	1005.743	871.489	41.316	39.882
64K	802.929	624.513	44.201	40.068

These stock mlx-lm benchmark arms were separated by other main-branch
commits, so they are supporting workload evidence rather than a strict
single-commit A/B. The optional 131K fused arm aborted with an
interactivity-impacting command-buffer error and is intentionally omitted.

A fixed dispatch threshold cannot pick correctly for all runtimes, so the
decision is exposed to the caller instead. This PR is the follow-up to the
maintainer request in #3658:

Let's add a force_fused=True option to mx.fast.scaled_dot_product_attention API that ignores the heuristics and uses fused path, and throws when a fused kernel is not available for the shape. — @zcbenz

Supersedes the closed #3660.
Addresses #3658.

API

mx.fast.scaled_dot_product_attention(
    q, k, v, *, scale, mask=None, sinks=None, force_fused=False, stream=None
)

• force_fused=False (default): preserves existing dispatch for previously
  supported shapes. The missing head_dim=192 vector kernel is added but
  reachable only through force_fused.
• force_fused=True: bypass the dispatch heuristics and use a fused kernel.
  Raises ValueError with the unsupported constraint when no fused kernel is
  available, which covers:
  • shapes not covered by the full- or vector-attention fused kernels
    (including unsupported head_dims, e.g. 512);
  • CPU execution;
  • training mode that requires a logsumexp output (the fused kernels do not
    emit one).

is_training on its own (when the VJP path already falls back) is allowed to
use the fused forward kernel — the output is identical.

Kernel support

• Full-attention (steel_attention): add head_dim=192/256 instantiations.
  Selected only when force_fused=True. Extend the existing large-head-dim
  V-tile synchronization barriers from BD == 128 to BD >= 128, covering
  192/256 as well as 128. This synchronization gap was also independently
  identified in giaki3003/mlx@1175b4d.
• Vector (decode) (sdpa_vector): add the missing head_dim=192 (192,192)
  and aggregation instantiations. The new 192 vector kernel is opt-in via
  force_fused like the new full-attention kernels; the existing head_dim=256
  vector kernel keeps its automatic routing.
• Route head_dim >= 192 away from the NAX kernel family, which has no
  192/256 instantiations.
• Check the compiled pipeline's maximum threadgroup size before both full
  attention dispatches, avoiding silent zero output on register-limited GPUs.
• Use each backend's actual support matrix: Metal full/vector kernels and CUDA
  cuDNN/vector kernels share the same opt-in semantics.

Validation

• Focused tests in python/tests/test_fast_sdpa.py:
  • force_fused=True full attention for head_dim 192/256 at short key
    lengths (where the default dispatch stays unfused) matches the reference;
  • force_fused=True full attention for head_dim 192/256 at kL=16385
    matches the reference, covering the Steel V-tile barrier path;
  • force_fused=True on an already-supported head_dim=128 matches;
  • force_fused=True on an unsupported shape (head_dim=512) raises;
  • unsupported full/vector head dims, causal qL > kL, vector GQA limits,
    and CPU execution raise with the specific violated constraint;
  • default dispatch with head_dim 192/256 full attention stays unfused and
    matches the reference;
  • head_dim=192 vector attention matches the reference.
• Existing test_fast_sdpa.py and test_fast.py suites continue to pass.

Scope

• Gemma 4 26B/31B are not affected: their long-context global layers use
  head_dim=512, which is outside this PR's scope (see Metal SDPA: no fused path for head_dim=512 (vector or full kernel), and the fallback is silent #3885).
  head_dim=256 layers there are bounded sliding-window layers.
• The Qwen 3.6 35B checkpoint measured in Allow force using fused kernel in scaled_dot_product_attention #3658 uses genuine
  head_dim=256 full attention and demonstrates the intended memory/throughput
  tradeoff. The optional 131K fused arm did not complete, so this PR does not
  claim a valid 131K A/B result.
• CUDA uses its cuDNN/vector support matrix for force_fused; a no-GPU build
  raises rather than silently ignoring the flag. CUDA compilation remains a
  GitHub CI validation boundary because the local validation host is macOS/Metal.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (API docstring)

[apocryphx]

Hi @hojin12312 — happy to report the M1 validation is done, and the results turned out more interesting than a simple thumbs-up. Thanks again for folding the safety and diagnostics work into your PR so cleanly; reviewing the capability-reason code while testing it on two machines was a pleasure, and I found nothing to object to — consider that portion reviewed and approved from my side.

The setup: your branch at 164fa7e23, plus a validation-only patch adding sdpa_vector 512/512 instantiations reachable through force_fused (the #3885 configuration — a test aid, not merge material). M1 Max 64 GB (applegpu_g13s) on macOS 26.5.2, with an M4 Max as control, both probing the byte-identical metallib.

1. The #3885 register cliff is still real — but it moved with the OS, and it's specialization-dependent. maxTotalThreadsPerThreadgroup for the d512 1-pass vector pipelines on the M1 Max today:

specialization	bf16	fp16
no mask	1024	832
causal	832	832
bool mask	832	832
float mask	832	832

(2-pass pipelines: 1024 in every variant; M4 Max: 1024 across the board.) Back in July, every d512 variant I probed reported 832. The current Metal compiler squeezes exactly one specialization under the dispatch size; everything else still caps below the kernel's hard-coded 1024-thread launch. I did not expect the ground to shift under this experiment in three weeks, but here we are.

2. The check behaves exactly as intended, end to end (force_fused, qL=1, kL=512 → 1-pass regime on 's'-class):

• bf16 no-mask: runs, matches the composite reference.
• fp16 no-mask: Maximum threads per threadgroup is 832 but requested 1024 for kernel sdpa_vector_float16_t_512_512_nomask_qnt_nc_nosinks. — a loud error precisely where the pre-Check threadgroup size in the 1-pass sdpa_vector dispatch #4018 code silently returned zeros.

One reproduction note for anyone retracing this: at kL >= 1024 on 's'-class devices the dispatcher selects the 2-pass kernels (which fit everywhere), so the capped 1-pass is only reachable at shorter kL. My first run sat exactly on the wrong side of that boundary and sailed through — worth knowing before anyone concludes the hazard is gone.

3. Your force_fused tests all pass on the M1 Max — the full set of seven — plus the whole test_fast_sdpa suite on the M4 Max (25 tests, one CI-memory skip).

The takeaway I'd offer for the PR description: the hazard these checks guard is not a static device property. It shifts with the Metal compiler version and with function-constant specialization — the same kernel is safe unmasked and unsafe masked, on the same chip, today. No support-matrix allowlist can track that, which makes the unconditional dispatch-time check exactly the right shape, and your full/NAX additions close the last unchecked dispatches.

I'll close #4186 now — glad the two efforts converged, and looking forward to seeing this land. The d512 vector instantiations can then follow as the small opt-in PR discussed in #3658, with their safety story now validated on real hardware.

[zcbenz]

Thanks!