Add d=512 vector SDPA kernels, reachable via force_fused

[apocryphx]

Proposed changes

Builds on #4185: adds vector-kernel (decode) coverage for head_dim = 512 —
the shape of Gemma 4's global attention layers (32 q / 4 kv heads) — reachable
only via force_fused. Automatic routing is unchanged, mirroring the
d=192 precedent from #4185: three-device sweeps
(#3885) show d=512 decode wins or loses depending on
device, shape, and depth, so the caller decides.

The 1-pass kernel runs D >= 512 with 16 simdgroups instead of 32. Register
pressure caps these pipelines below the 1024-thread launch on some GPUs, and
the cap is not a static device property: on the same M1 Max we have observed
three different cap configurations across five weeks (all-832 in July; 1-pass
bf16-nomask at 1024 with everything else 832 in early August; nearly
everything 832 including most 2-pass specializations by mid-August). A
portable dispatch has to fit the floor — 512 threads does, under every cap
observed — and if a future compiler squeezes below that, the dispatch-time
check from #4018 turns it into a loud error rather than silent zeros.
Lane-guarded reductions and per-simdgroup output slices keep the math
identical to the BN=32 form.

Validation

• M1 Max (applegpu_g13s) + M4 Max (applegpu_g16s): force_fused output
  matches the composite path to input-dtype rounding for bf16 and fp16 across
  1-pass and 2-pass regimes, MHA/GQA/MQA shapes; the d=256 control is
  bit-exact (force and auto select the same kernel there).
• Throw paths verified on both machines: d=512 prefill (no full kernel) and
  d=72 (no kernel anywhere) raise with the supported-dims message; the vector
  message now lists 512.
• python -m pytest python/tests/test_fast_sdpa.py: 19 passed, 2 skipped,
  705 subtests — including new d=512 vector-loop and short-kL 1-pass cases
  (fp16, masked and unmasked — the specializations M1-class caps hardest).
  Rebased onto main post-Add force_fused option to scaled_dot_product_attention #4185 merge; full test_fast_sdpa.py and
  test_fast.py pass unchanged on the new base.
• Decode perf landscape for the routing decision: three-device sweep data
  (M4 Max, M1 Max, M4 iPad Pro) in Metal SDPA: no fused path for head_dim=512 (vector or full kernel), and the fallback is silent #3885.

🤖 Generated with Claude Code

[nastya236]

Thanks for your contribution! However, we have a new policy regarding AI-generated code (see #4331), so I’m closing this pull request. Feel free to open a new one that follows the updated policy.