Fit FP32 wide-head attention tiles within Metal shared memory

[Gajesh2007]

Summary

Forced full attention with FP32 head dimensions 192 or 256 can exceed Metal's 32 KiB threadgroup-memory limit. On M4 Max, the D256 bq32/bk16/wm4 specialization requires 53,760 bytes and fails to load.

Use bq16/wm2, with bk16 for D192 and bk8 for D256. The padded shared arrays require 27,904 and 28,928 bytes respectively. Computation stays FP32; existing default-routing heuristics and other dtypes are unchanged.

Before

flowchart LR
  A[Forced FP32 full attention] --> B[sdpa_full_self_attention: bq32 bk16 wm4]
  B --> C[Wide-head shared arrays exceed 32 KiB]
  C --> D[Metal rejects kernel]

Loading

After

flowchart LR
  A[Forced FP32 full attention] --> B{Head dimension}
  B -->|192| C[bq16 bk16 wm2]
  B -->|256| D[bq16 bk8 wm2]
  B -->|other| E[Existing tile]
  C --> F[Fused attention result]
  D --> F
  E --> F

Loading

Validation

Apple M4 Max, Swift 6.3, using the new C/Swift explicit-fused binding: five tests passed, covering default-route compatibility, pitched BOOL masks and sinks against a dense reference for D64/D192/D256, wide-head GQA8, unsupported D512 refusal, and an allocation check below 16 MiB where a composed score tensor alone would occupy 32 MiB.

The companion packed-KV integration also passed 40 tests, including exact Qwen and GPT-OSS attention geometries. Performance qualification is in progress in Layr-Labs/d-inference#860; this patch fixes kernel eligibility and does not claim a throughput improvement.