Commit Graph

21 Commits

Author SHA1 Message Date
modelbeast
e8f1721996 docs: remove the now-refuted width framing; intro says six machines 2026-07-17 14:03:44 +10:00
modelbeast
5f87aec717 docs: M1 Max (6th machine) refutes the width theory — report the correction
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.
2026-07-17 14:02:54 +10:00
modelbeast
6b651f6f8e docs: intro reflects 5 machines 2026-07-17 01:09:00 +10:00
modelbeast
6f0fc4998b docs: M3 Air — matched-width generational delta + fanless negative result
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.
2026-07-17 01:08:19 +10:00
modelbeast
64be04aeac docs: retitle for the 4-machine matrix (M1/M3/M4/M5, Ultra->laptop)
Also validated the tiled-compile patch against the repo's own standards:
ruff check + ruff format clean, pytest 64 passed / 0 failed.
2026-07-16 23:56:38 +10:00
modelbeast
fd00c31b1d docs: add M5 — head_dim padding is worth ~5x on laptops (vs ~2x on Ultras)
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.
2026-07-16 23:54:49 +10:00
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
modelbeast
a9f6b75ef2 docs: M-series fleet ablation (M3 Ultra / M1 Ultra / M4 Pro 24GB)
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).
2026-07-16 22:48:06 +10:00
cmoyates
e91089a804
docs: wave2 ablation benchmarks + optimization plan + brainstorm
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>
2026-03-09 19:32:04 -02:30
cmoyates
f73aabf5af
chore: default bf16+fused decode on, add benchmark results to plan
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>
2026-03-08 21:27:56 -02:30
cmoyates
b3b7df271d
feat: MLX memory optimizations — bf16, fused decode, deterministic GC
Phase 1: selective bf16 mixed precision (backbone fp32, decoders bf16,
sigmoid always fp32). dtype param on GreenFormer, all outputs guaranteed
fp32.

Phase 2: FusedDecoderPair batches 3 upsample ops by concatenating
alpha+fg projections along channel axis. Bit-exact parity with unfused
path.

Phase 3: deterministic GC pipeline in tile loop (del + gc.collect +
mx.clear_cache) and engine.py intermediate tensor cleanup.

7 new tests: bf16 dtype/range/nan, fused-vs-unfused parity, fp32
default unchanged. 80 passed, 0 failures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:21:01 -02:30
cmoyates
830e88f050
docs: add brainstorm + plan for MLX memory optimizations
bf16 mixed precision, batched decoder upsample, deterministic GC pipeline.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-08 21:12:59 -02:30
cmoyates
88a0cd9cb8
docs: add deep modules refactor plan
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-03 09:45:43 -03:30
cmoyates
8be8d003e7
feat: add 2048 smoke test for native-resolution validation
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>
2026-03-01 07:04:28 -03:30
cmoyates
4bad586112
feat: add CorridorKeyMLXEngine integration surface
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>
2026-03-01 06:43:07 -03:30
cmoyates
57c3aff913
feat(phase6): optimization, benchmarking, and tiled inference
- 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>
2026-03-01 06:24:10 -03:30
cmoyates
0481955cfd
feat(phase5): full model assembly + e2e parity
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>
2026-03-01 06:07:43 -03:30
Cristopher Yates
c2e7967538
feat(phase4): MLX Hiera backbone port (#2)
* feat(phase4a): MLX Hiera backbone with unroll/reroll parity

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>

* feat(phase4c): weight loading + pos_embed bicubic interpolation

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>

* feat(phase4d): shape + parity tests, fix attention transpose

Fix MaskUnitAttention output transpose (0,2,3,1,4 → 0,3,2,1,4) to
match PyTorch transpose(1,3) token ordering for windowed attention.

Parity results (4 stages):
  Stage 0: max_abs 2.9e-4, mean 1.6e-5
  Stage 1: max_abs 1.3e-4, mean 1.3e-5
  Stage 2: max_abs 1.1e-2, mean 5.2e-5 (16 blocks, expected drift)
  Stage 3: max_abs 5.8e-4, mean 2.6e-5

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* docs: clarify pos_embed parameter intent in HieraBackbone

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:55:24 -03:30
Cristopher Yates
150c974eed
feat(phase3): PyTorch→MLX weight converter (#1)
* feat(phase3): PyTorch→MLX weight converter + safetensors output

365 keys mapped (367 source - 2 num_batches_tracked skipped).
Conv weights transposed (O,I,H,W→O,H,W,I), refiner stem remapped,
4ch patch embed preserved. 12 conversion tests + diagnostic report.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix(phase3): conv allowlist, unused var, test fixture

- Replace _is_conv_weight heuristic with explicit CONV_WEIGHT_KEYS frozenset (15 keys)
- Remove unused `skipped` list from convert_state_dict
- Module-scoped pytest fixture eliminates 11 redundant checkpoint loads

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:16:56 -03:30
cmoyates
091261ccff
feat(phase2): MLX decoder + refiner with parity tests
- MLP, DecoderHead (SegFormer-style), RefinerBlock, CNNRefinerModule
- All NHWC native, GroupNorm with pytorch_compatible=True
- Layout utils: NCHW<->NHWC, conv weight transpose (OIHW->OHWI)
- Dump script extended to export decoder/refiner weights
- Parity: decoder max_abs ~5e-6, refiner max_abs ~1e-5

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:00:00 -03:30
cmoyates
897242d200
init: repo scaffolding, deps, prompts, plan
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 04:48:51 -03:30