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>
HieraBackbone.load_checkpoint() loads safetensors, strips encoder.model.
prefix, bicubic-interpolates pos_embed from 512x512 to 128x128 tokens.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Port timm hiera_base_plus_224 to MLX: PatchEmbed, MaskUnitAttention,
HieraBlock, unroll/reroll, HieraBackbone. 4 NHWC feature maps at
strides 4/8/16/32. backbone.py now re-exports from hiera.py.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
