Completes every Apple Silicon GPU generation ever shipped (M1 2020 -> M5).
Headline: tiled 768/64+compile keys a full 2048 frame in ~2.2GB on EVERY
machine with IDENTICAL alpha (MAE 0.00849) — M3 Ultra 1949ms ... M1 mini 8GB
14269ms. Full-frame needs 26GB and is unreachable for most Macs ever sold, so
tiled should arguably be the default, not the fallback. The README's 'might
even work on your old MacBook Pro' undersells it.
Also: real CorridorKey ablation on M1 Max (stage_gc harmful at 32GB too:
0.81x @512), and the OG 8-core M1 mini confirms the hdim-56 fallback (1.00x
x3, padding 1.9x) — now 8/8 machines.
hdim-56 fallback is textbook here (1.00x x3); padding = 2.0x. Fits the
corrected fused-kernel-quality theory as a prediction, not a fit: 0.50x
fused/unfused -> 2.0x padding win, exactly between M3 Air (0.40x->2.6x) and
M1 Max (0.87x->1.3x). At 10c matched width the ladder is 138.9/48.8/44.0ms at
hdim56 but 54.3/24.4/8.9ms at hdim64 — M5's gains live almost entirely in the
fused kernel, making the head-dim fix most valuable on the newest silicon.
Padding 56->64 helps everywhere (1.3x-4.9x) and the hdim-56 silent fallback
reproduces on all six machines. But the M1 Max (24c) gains LEAST of all,
breaking our 'narrower GPU => bigger win' reading (M5 10c gained most).
Real predictor: fused-kernel quality vs that chip's own matmul throughput —
M1 Max's fused path is only 0.87x unfused at every size, so missing it costs
little; M5's is 0.14x, so missing it costs a lot. Corrected in place rather
than left as a tidy-but-wrong story.
Fifth machine. Two additions:
- M3 Air (10c) vs M5 (10c): same GPU width, 2 generations apart -> M5 is 7.7x
at seq 4096. Isolates architecture from width (the Ultras can't). head_dim-56
fallback reproduces here too (padding = 2.6x), so it now holds on every
Apple GPU generation shipping today.
- Fanless sustained load: NO meaningful throttle over 5 min (1.01x at 300s).
We predicted a decay curve and found none — reporting the negative, since it
validates that the burst numbers aren't short-run artifacts.
M5 MacBook Pro (10-core GPU, 16GB) reproduces the head_dim=56 silent fallback
and shows the fix matters far MORE on narrow GPUs: padding 56->64 gives 4.9x
(1,16,4096), 5.1x (1,16,2048), 3.2x (4096,2,64). Raw fused-vs-unfused at
4096/fp32 is 0.14x on M5 vs ~0.53x on M3 Ultra — an M5 laptop matches an M1
Ultra on fused attention but is ~4x slower unfused. This reframes the fix as a
laptop optimization, not a big-iron one.
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.
Ran scripts/bench_optimizations.py unmodified across three machines + an
engine-level tiled×compile extension with ground-truth alpha scoring.
Headlines: sdpa = 4.5x regression on M1 at 2048 (bf16 innocent — isolated
per-toggle on both Ultras); stage_gc harmful on Ultra-class; tiled 768/64
+ compile = best 2048 config on both Ultras; 24GB Macs swap-bound at 2048
full-frame (tiled required).
Tiles are fixed-shape, so fused compilation applies the same as full-frame.
Forcing compile=False in the tiled branch left performance on the table:
measured on the same 12-frame 2048px green-screen set (alpha output
bit-identical, IoU 0.9316):
M3 Ultra: 3.64 -> 2.48 s/frame (1.47x)
M1 Ultra: 4.65 -> 4.19 s/frame (1.11x)
Peak memory unchanged (~2.3 GB).
Ablation sweep across 512/1024/2048 resolutions with 8 configs.
Tiled 768/64 = 1.5x faster + 11.4x less memory vs full-frame at 2048.
stage_gc net negative at small res, breaks even at 2048.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
PIL bicubic vs MLX cubic diff too large (0.22-0.59) — resize stays
on CPU. Only normalize+concat moved to MLX. Tiled path unchanged.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Fold window dim into batch (transpose before reshape), call SDPA,
unfold back. All parity tests pass with 0 regression.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
mx.eval + gc + clear_cache at backbone→decoder and decoder→refiner
boundaries in eager mode. _compiled flag skips GC under mx.compile.
Compilation tests updated to use compile_model() API.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
slim=True returns 4-key dict, engine uses it, pipeline.infer keeps full.
2 new tests: key count + bit-exact match vs full output.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
tile_size/overlap params on CorridorKeyMLXEngine — model loads at
tile_size, input stays full-res, forward uses tiled_inference w/ GC.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
load_model() now defaults to dtype=bf16, fused_decode=True. Both are
free (zero parity regression, bit-exact fused path). Backbone/sigmoid
stay fp32.
Plan updated with benchmark results: tiled+GC = 12x peak memory
reduction at 2048x2048 (27.6GB → 2.3GB), all acceptance criteria met.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
PyTorch reference script had wrong concat order (forward) vs the
original nikopueringer/CorridorKey which uses torch.cat([c4,c3,c2,c1]).
Golden fixtures regenerated. Relax 3 tight e2e tolerances for Metal
vs CPU float32 drift.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace mypy with Astral's ty for faster type checking. Fix
Image.BICUBIC deprecation and inline **kwargs for ty compatibility.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Torch decoder concatenates feature projections as [c4,c3,c2,c1] but
MLX was using [c1,c2,c3,c4]. The linear_fuse conv was trained with
the Torch order, producing scrambled features. Verified fix gives
correlation=1.0 against Torch on identical inputs.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Script + pytest tests to verify MLX inference works at CorridorKey's
training resolution. Uses samples/ by default, synthetic fallback.
Reports timing, peak memory, output diagnostics.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
preprocess() already calls normalize_rgb() internally; calling it
before passing to preprocess() applied ImageNet normalization twice.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Drop-in MLX backend for main CorridorKey repo. Engine class wraps
existing model/inference with process_frame() API returning alpha,
fg, comp, processed as uint8 numpy arrays. Lowers Python to >=3.11,
removes unused deps, adds smoke script and 16 contract tests.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Cache nn.Upsample instances in DecoderHead/GreenFormer __init__
(eliminated ~7 allocations per forward pass)
- Add mx.compile() support via load_model(compile=True)
- Benchmark harness: eager vs compiled, multi-resolution, parity checks
- Tiled inference with overlap blending for large images
- Profiling utilities with forced mx.eval for accurate timing
- Reference comparison script (scripts/compare_reference.py)
- 12 new tests (compiled consistency + tiling)
- README performance section with Apple Silicon guidance
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Wire GreenFormer (backbone + decoders + refiner), add image I/O,
inference pipeline, CLI entry point, and 3 test files (23 new tests).
E2e parity vs PyTorch: max abs err < 1.5e-4 across all outputs.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
