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:
parent
a9f6b75ef2
commit
df28940d81
@ -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 1–5 % peak-memory saving. Wave2 measured
|
||||
|
||||
Loading…
Reference in New Issue
Block a user