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