Add force_fused option to scaled_dot_product_attention

[apocryphx]

Implements the option proposed by @zcbenz in #3658 (#3658 (comment)), with the semantics posted there:

• force_fused=True bypasses the use_fallback routing heuristics and always uses a fused kernel; whether the vector or the full kernel runs still follows query length, as with default routing.
• When no fused kernel supports the configuration, it throws a descriptive error with the offending shapes and the supported-dims list — loud, never silent. Example:

  [scaled_dot_product_attention] force_fused=true but no fused kernel is available: the full attention kernel supports head dims {64, 80, 96, 128} with matching query/value head dims; got query head dim 512 and value head dim 512.

• Default false; no behavior change for existing callers.
• CUDA consults the same flag against its own support matrix (cuDNN + vector), so semantics stay uniform across backends; CPU streams and no-GPU builds report a clear error.

Implementation notes:

• The capability half of use_fallback moves into a new per-backend static fused_unsupported_reason (Metal, CUDA, no_gpu), so the routing predicate and the force_fused error report share one source of truth rather than duplicating the head-dim lists. use_fallback itself is now the heuristics (e.g. prefer unfused while training) plus that capability check — behavior unchanged.
• Safety pairing: adds check_kernel_threadgroup_size to the two full-attention dispatches (steel and NAX) — after Check threadgroup size in the 1-pass sdpa_vector dispatch #4018 and Fix 2pass sdpa on < M2 #3099 these were the only SDPA dispatches still launching unchecked. force_fused makes register-limited configurations reachable, and on M1-class GPUs an oversized launch is silently dropped and returns zeros (measured in Metal SDPA: no fused path for head_dim=512 (vector or full kernel), and the fallback is silent #3885).
• On Metal, force_fused also works under gradient tracing (fused forward, gradients by recomputation through the reference implementation) since the Metal SDPA VJP fallback is unconditional; the logsumexp capability guard exists for backends with a fast VJP.

Tests (Metal, M4 Max): new TestSDPAForceFused — fused-vs-reference match on the vector and full paths (with GQA and causal masks), per-constraint error cases (head dims for each kernel, qL×GQA bound, causal with qL>kL), and the CPU-stream error. Full test_fast_sdpa and test_fast suites pass with no default-path changes.

🤖 Generated with Claude Code

[hojin12312]

Hi @apocryphx — looks like we ended up working on this in parallel. Thanks for catching the unchecked full-attention launches and for spelling out the capability errors so clearly.

I had opened #4185 earlier with the 192/256 full-attention kernels that motivated #3658. I’ve since updated it with the full/NAX threadgroup checks and per-backend unsupported diagnostics, including CUDA, while preserving the existing C++ positional stream argument. I’ve also credited your work in the PR description.

Since the two PRs now overlap quite a bit, would you be open to converging on #4185? It now combines the 192/256 kernel support with the safety and diagnostics from your approach. If there’s anything from #4186 I’ve missed, I’d be happy to incorporate it.

[apocryphx]

Happy to consolidate on #4185 — the 192/256 kernels plus the shared safety and diagnostics work belong in one PR, and @hojin12312 keeping the C++ positional stream argument is the better API call (this branch inserted force_fused before s, which would have broken existing positional C++ callers).

On the two asks: I'm reviewing the capability-diagnostics portion of #4185 now. For M1-class validation I'm re-running the #3885 reproduction against the #4185 branch — pipeline probe of the register-limited d512 configurations plus confirming an oversized launch raises the threadgroup-size error instead of silently returning zeros — and will post results in #4185 once the M1 run completes. I'll close this PR at that point.

[apocryphx]

Consolidated into #4185 — M1 validation posted there. Thanks all.