docs: root-cause the sdpa finding — head_dim=56 misses the fast path

Shape-spy on a real 2048 forward + exact-shape sweeps on M1/M3 Ultra:
fused==unfused at hdim 56 (silent fallback) on both machines, so use_sdpa
buys nothing and its transpose choreography costs ~4.5x on M1. Padding to
hdim 64 engages the fast kernel: 2.2x (M3) / 2.7x (M1) on the dominant
global-attention shape despite 14% extra FLOPs.
This commit is contained in:
modelbeast 2026-07-16 23:05:13 +10:00
parent a9f6b75ef2
commit df28940d81

View File

@ -24,10 +24,31 @@ Apple Silicon range. Contributed from a 3-Mac render-farm setup; happy to run fo
| bf16+sdpa | 6388 ms | 1767 ms |
The suspected M1 danger (bf16 emulation — M1 lacks hardware bf16) is **innocent** on this
workload; `mx.fast.scaled_dot_product_attention` at 2048-scale sequence lengths is the
cliff. bf16 partially mitigates it (fp32 sdpa is the worst path). Suggestion: gate
`use_sdpa` by GPU generation rather than defaulting it on. (This may interest MLX
upstream too — it looks like a Metal kernel path issue on M1, not a model issue.)
workload.
**Root cause (follow-up, same day):** it is *not* the SDPA kernel. Instrumenting a real
2048 forward shows the model's attention shapes are `head_dim=56` (global blocks
`(1,8,16384,56)` ×15, `(1,16,4096,56)` ×3; windowed `(4096,·,·,56)`). At head_dim 56,
`mx.fast.scaled_dot_product_attention` **silently falls back to the unfused path**
fused == unfused within 1 % on every shape on both M1 and M3 Ultra. So `use_sdpa=True`
currently buys zero kernel benefit on any machine, while its branch pays extra 5-D
transpose/reshape choreography — which is what costs ~4.5× on M1 (layout-sensitive),
~free on M3.
**The actionable win: pad head_dim 56 → 64.** The fast path then engages and beats
unfused decisively — *despite ~14 % more FLOPs* (fp32, exact model shapes):
| shape (hdim 64) | M3 fused/unfused | M1 fused/unfused |
|---|---|---|
| (1, 8, 16384, ·) | 28.6 / 63.1 ms → **2.2×** | 53.8 / 144.5 ms → **2.7×** |
| (1, 16, 4096, ·) | 3.8 / 7.2 ms → 1.9× | 7.5 / 11.1 ms → 1.5× |
| (4096, 2, 64, ·) | 1.3 / 2.5 ms → 1.9× | 1.5 / 2.6 ms → 1.7× |
Suggested changes: (a) pad qkv projections to head_dim 64 (at minimum for the global
blocks) so sdpa's fast kernel actually engages — M1 benefits *more* than M3; (b) until
then, default `use_sdpa=False` (it is currently overhead-only). For MLX upstream: a
warning (or doc note) when sdpa silently falls back on unsupported head dims would have
made this obvious much sooner.
2. **`stage_gc` is harmful on Ultra-class machines at every resolution** — 0.53× at 512,
0.64× at 1024, 0.66× at 2048 on M3 Ultra, for a 15 % peak-memory saving. Wave2 measured