From df28940d81226adfa8d33482b9cf4791989f4086 Mon Sep 17 00:00:00 2001 From: modelbeast Date: Thu, 16 Jul 2026 23:05:13 +1000 Subject: [PATCH] =?UTF-8?q?docs:=20root-cause=20the=20sdpa=20finding=20?= =?UTF-8?q?=E2=80=94=20head=5Fdim=3D56=20misses=20the=20fast=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- ...6-07-16-m-series-fleet-ablation-results.md | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/docs/2026-07-16-m-series-fleet-ablation-results.md b/docs/2026-07-16-m-series-fleet-ablation-results.md index 87295da..48ea661 100644 --- a/docs/2026-07-16-m-series-fleet-ablation-results.md +++ b/docs/2026-07-16-m-series-fleet-ablation-results.md @@ -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