monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
df28940d81
corridorkey-mrp-mlx/docs/2026-07-16-m-series-fleet-ablation-results.md
modelbeast df28940d81 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.
2026-07-16 23:05:13 +10:00

7.2 KiB
Raw Blame History

title date device script
M-series fleet ablation: M3 Ultra / M1 Ultra / M4 Pro 2026-07-16 M3 Ultra 256GB · M1 Ultra 128GB (64-core GPU) · M4 Pro Mac mini 24GB scripts/bench_optimizations.py (unmodified) + engine-level tiled×compile extension

M-series Fleet Ablation Benchmarks

Same methodology as the 2026-03-09 wave2 doc (--sweep ablation --resolution 512 1024 2048, checkpoint weights, 3 warmup + 10 bench runs), executed on three machines spanning the Apple Silicon range. Contributed from a 3-Mac render-farm setup; happy to run follow-ups.

Headline findings

  1. sdpa is a ~4.5× regression on M1-class GPUs at 2048 — and free on M3. Single-toggle isolation at 2048 (median of 6, all-off baseline):

    toggle M1 Ultra M3 Ultra
    all off 4139 ms 1743 ms
    bf16 only 4120 ms (neutral!) 1766 ms
    sdpa only 18592 ms (4.5× slower) 1763 ms
    bf16+sdpa 6388 ms 1767 ms

    The suspected M1 danger (bf16 emulation — M1 lacks hardware bf16) is innocent on this 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 a mild 0.79–0.99× on its reference hardware; on big-memory machines it's pure overhead. Suggestion: document as a low-memory-only flag.

  3. Tiled + compile is the best 2048 config on both Ultras — upstream engine.py forced compile=False in tiled mode; tiles are fixed-shape, so fused compilation applies (patched in this fork, output bit-identical, verified against ground-truth alpha):

    engine config (2048 input) M3 Ultra M1 Ultra peak alpha MAE*
    full-frame 2048 2788 ms 5373 ms 27.9 GB 0.00906
    tiled 512/64 2760 ms 4576 ms 2.2 GB 0.00821
    tiled 512/64 + compile 2478 ms 4222 ms 2.3 GB 0.00821
    tiled 768/64 + compile 1949 ms 3275 ms 2.4 GB 0.00842
    tiled 1024/64 + compile 3475 ms 5864 ms 3.7 GB 0.00915

    *alpha MAE vs exact ground truth: synthetic 2048² green-screen plates (soft-alpha subject + motion-blur stripes + defocus disk composited over chroma green), hint = 8× downscaled truth. Tiled beats full-frame on accuracy as well as memory — the model runs at native tile scale over full-res input. Confirms wave2's tiled-768 pick and adds ~8–13 % from compiling the tile graph.

  4. 24 GB Macs must tile at 2048. On the M4 Pro mini (24 GB), every full-frame 2048 config lands at 20–24 s/run — the ~26 GB working set swaps; toggle choice becomes noise. At ≤1024 the M4 is healthy (218 ms @512, 1149 ms @1024). Tiled 2048 runs in ~2.3 GB.

Cross-machine baselines (all-off)

res M3 Ultra M1 Ultra M4 Pro 24GB wave2 reference
512 53.3 ms 80.8 ms 218.1 ms 119.6 ms
1024 247.0 ms 350.6 ms 1149.2 ms 610.7 ms
2048 1750.2 ms 3527.3 ms 23147 ms (swap) 4984.7 ms

Full ablation tables

M3 Ultra 256GB

Config 512 1024 2048 peak @2048
baseline 53.3 247.0 1750.2 26689 MB
slim+sdpa+bf16+fused_decode+gpu_preprocess 53.5 248.2 1790.1 27245 MB
slim+stage_gc+bf16+fused_decode+gpu_preprocess 95.4 382.7 2644.3 26689 MB
slim+stage_gc+sdpa+fused_decode+gpu_preprocess 96.3 385.9 2676.9 26661 MB
slim+stage_gc+sdpa+bf16+gpu_preprocess 99.4 372.1 2618.6 26661 MB
slim+stage_gc+sdpa+bf16+fused_decode 99.0 386.5 2688.0 26661 MB
stage_gc+sdpa+bf16+fused_decode+gpu_preprocess 100.1 391.4 2628.2 26661 MB
slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess 101.0 390.9 2617.9 26661 MB

M1 Ultra 128GB

Config 512 1024 2048 peak @2048
baseline 80.8 350.6 3527.3 26689 MB
slim+sdpa+bf16+fused_decode+gpu_preprocess 79.1 353.3 33137.8 27245 MB
slim+stage_gc+bf16+fused_decode+gpu_preprocess 189.5 705.5 4919.0 26689 MB
slim+stage_gc+sdpa+fused_decode+gpu_preprocess 195.6 711.1 4939.2 26661 MB
slim+stage_gc+sdpa+bf16+gpu_preprocess 198.2 681.6 4853.8 26661 MB
slim+stage_gc+sdpa+bf16+fused_decode 197.7 714.1 4822.9 26661 MB
stage_gc+sdpa+bf16+fused_decode+gpu_preprocess 194.2 733.2 6274.1 26661 MB
slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess 193.8 737.1 5128.2 26661 MB

Note the 33.1 s outlier: the only 2048 config without stage_gc but with sdpa — sdpa's slow path dominating once nothing throttles it (see headline 1 for the isolation). The sdpa cliff appears only at 2048; 512/1024 are unaffected.

M4 Pro Mac mini 24GB

Config 512 1024 2048 (swap-bound)
baseline 218.1 1149.2 23147.1
slim+sdpa+bf16+fused_decode+gpu_preprocess 221.4 1139.0 20994.5
slim+stage_gc+… (all stage_gc combos) 266–273 1268–1292 20233–24253

At 2048 the ~26 GB working set exceeds 24 GB unified memory; all configs swap and differences are not meaningful. ≤1024 full-frame or tiled-anything is the usable envelope.

Recommended settings by hardware

Hardware 2048 recommendation
M3-class (Max/Ultra) full-frame or tiled 768/64 + compile; every toggle optional; avoid stage_gc
M1/M2-class avoid sdpa; tiled 768/64 + compile; avoid stage_gc
≤24–32 GB any gen tiled required (2.3 GB vs 26 GB); tiled 768/64 + compile

Environment: MLX (venv per uv sync --extra mlx), macOS 26.5, checkpoint v1.0.0. Quality harness (ground-truth plates + scoring) available on request — it's ~150 lines and reproduces the alpha-MAE column.