monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: wave2 ablation benchmarks + optimization plan + brainstorm

Browse Source 
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>
This commit is contained in:
cmoyates 2026-03-09 19:32:04 -02:30
parent d13ba1639d
commit e91089a804
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
4 changed files with 838 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

281
docs/MLX Vision Transformer Optimization Techniques.md Normal file
View File

File diff suppressed because one or more lines are too long

83
docs/benchmarks/2026-03-09-wave2-ablation-results.md Normal file
View File

@ -0,0 +1,83 @@
---
title: "Wave 2 optimization ablation benchmarks"
date: 2026-03-09
device: "Apple Silicon (unified memory)"
script: scripts/bench_optimizations.py
---
# Wave 2 Optimization Ablation Benchmarks
## Ablation Sweep (baseline = all optimizations off)
Toggle flags: `slim`, `stage_gc`, `sdpa`, `bf16`, `fused_decode`, `gpu_preprocess`
### 512x512
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|---|---|---|---|---|
| baseline | 119.6 | 118.3 | 2419 | 1.00x |
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 120.0 | 119.2 | 2413 | 1.00x |
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 147.8 | 146.7 | 2306 | 0.81x |
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 149.2 | 147.9 | 2344 | 0.80x |
| slim+stage_gc+sdpa+bf16+fused_decode | 149.3 | 147.6 | 2344 | 0.80x |
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 150.5 | 148.7 | 2344 | 0.79x |
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 151.6 | 148.5 | 2344 | 0.79x |
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 151.9 | 149.1 | 2344 | 0.79x |
### 1024x1024
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|---|---|---|---|---|
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 583.3 | 579.9 | 3819 | 1.05x |
| baseline | 610.7 | 606.2 | 3673 | 1.00x |
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 655.5 | 650.6 | 3673 | 0.93x |
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 655.9 | 650.6 | 3673 | 0.93x |
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 657.5 | 650.6 | 3673 | 0.93x |
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 675.3 | 667.7 | 3673 | 0.90x |
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 686.0 | 682.7 | 3673 | 0.89x |
| slim+stage_gc+sdpa+bf16+fused_decode | 687.5 | 680.4 | 3673 | 0.89x |
### 2048x2048
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | vs Baseline |
|---|---|---|---|---|
| baseline | 4984.7 | 4954.6 | 26689 | 1.00x |
| stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 5016.8 | 4840.4 | 26661 | 0.99x |
| slim+stage_gc+sdpa+bf16+fused_decode+gpu_preprocess | 5048.9 | 4848.1 | 26661 | 0.99x |
| slim+sdpa+bf16+fused_decode+gpu_preprocess | 5175.5 | 5065.2 | 27245 | 0.96x |
| slim+stage_gc+sdpa+bf16+fused_decode | 5375.2 | 5083.1 | 26661 | 0.93x |
| slim+stage_gc+sdpa+bf16+gpu_preprocess | 5396.3 | 5357.0 | 26661 | 0.92x |
| slim+stage_gc+sdpa+fused_decode+gpu_preprocess | 5683.3 | 5546.8 | 26661 | 0.88x |
| slim+stage_gc+bf16+fused_decode+gpu_preprocess | 5771.3 | 5643.0 | 26689 | 0.86x |
## Tiled vs Full-Frame at 2048x2048
| Config | Median (ms) | Min (ms) | Peak Mem (MB) | Speed vs FF | Mem vs FF |
|---|---|---|---|---|---|
| full-frame 2048 | 5031 | 4838 | 26281 | 1.0x | 1.0x |
| **tiled 768/64** | **3344** | **3311** | **2302** | **1.5x** | **11.4x less** |
| tiled 512/64 | 3978 | 3963 | 2133 | 1.3x | 12.3x less |
| tiled 512/128 | 4055 | 4025 | 2133 | 1.2x | 12.3x less |
| tiled 1024/64 | 6144 | 6125 | 3425 | 0.8x | 7.7x less |
## Key Findings
1. **stage_gc hurts at 512/1024, breaks even at 2048.** GC pauses cost more than memory savings at small resolutions. Only at 2048 does the computation graph get large enough to justify intermediate materialization.
2. **No single optimization consistently beats baseline across resolutions.** At 512, baseline wins. At 1024, slim+sdpa (no stage_gc) wins by 5%. At 2048, nothing reliably beats baseline.
3. **Memory scales super-linearly: 2.4GB → 3.7GB → 26.7GB** for 512 → 1024 → 2048. The 7x jump from 1024→2048 indicates lazy graph accumulation across 24 backbone blocks dominates at high resolution.
4. **Tiled 768/64 is the clear winner at 2048** — 1.5x faster AND 11.4x less memory than full-frame. The smaller per-tile computation graph avoids the massive intermediate state buildup.
5. **bf16 can hurt at 2048** — possible promotion overhead in large graphs.
6. **gpu_preprocess is the most impactful single flag** — dropping it consistently hurts latency.
7. **Phase 6 (GPU tile accumulators) is low priority** — numpy accumulator transfer overhead is negligible vs compute. The memory problem is entirely graph-side.
## Recommendations
- Default to **tiled 768/64** for 2048x2048 production use
- For ≤1024, use **full-frame with slim+sdpa+bf16+fused_decode+gpu_preprocess**
- stage_gc should be **off by default**, enabled only if memory-constrained at ≥2048

143
docs/brainstorms/2026-03-09-pytorch-optimization-learnings-brainstorm.md Normal file
View File

@ -0,0 +1,143 @@
# PyTorch Optimization Learnings for MLX Port
**Date:** 2026-03-09
**Status:** Draft
**Sources:**
- [PR #104](https://github.com/nikopueringer/CorridorKey/pull/104) (MarcelLieb) — fp16/compile/GPU preprocessing
- [CorridorKey_Test](https://github.com/Raiden129/CorridorKey_Test) (Raiden129) — FlashAttention/tiled refiner/token routing/cache clearing
- Deep research: MLX Vision Transformer Optimization Techniques (local)
## What We're Exploring
Extract optimization techniques from PyTorch CorridorKey forks, assess MLX applicability, identify net-new improvements for our port.
## Source Analysis
### PR #104 (MarcelLieb) — VRAM: 18 GB -> 1.9 GB
| Technique | Description | MLX Status |
|---|---|---|
| `torch.compile()` on GreenFormer | Graph-level JIT fusion | **Done** — `mx.compile()` in pipeline.py |
| Full fp16 model weights | `model.to(model_precision)` | **Done** — bf16 decoders, fp32 backbone/refiner |
| Mixed precision toggle | Conditional `torch.autocast` | **Done** — per-component dtype in GreenFormer |
| `torch.inference_mode()` | Faster than `no_grad()` | **N/A** — MLX has no grad tracking in inference |
| `set_float32_matmul_precision("high")` | TF32 matmuls | **N/A** — Apple Silicon different numerics |
| GPU-side preprocessing | Moved normalize/resize from numpy/cv2 to torch tensors on device | **NEW** — worth investigating |
### Raiden129/CorridorKey_Test — VRAM: 9.8 GB -> 1.6 GB (84% reduction), 40% faster
| Technique | Description | MLX Status |
|---|---|---|
| FlashAttention patching | Squeeze 5D->4D Q/K/V for SDPA dispatch | **Different** — MLX has own heuristic (see below) |
| Tiled CNN refiner | 512x512 tiles, 128px overlap, linear blend | **Done** — tiling.py (but uses numpy accumulators) |
| cuDNN benchmark disable | Avoid workspace allocation | **N/A** — no cuDNN on Metal |
| Strategic cache clearing | `empty_cache()` between encoder/decoder/refiner | **Partial** — done in tiling, NOT in non-tiled path |
| Token routing (experimental) | Route easy tokens to LTRM, edge tokens to full attention | **NEW** — novel compute reduction |
## Net-New Opportunities
### 1. Verify MLX Attention Memory Behavior
**Problem:** MLX's `mx.fast.scaled_dot_product_attention` has a dynamic heuristic choosing between O(N) streaming (flash-like) and O(N^2) explicit materialization. We don't know which path our Hiera backbone triggers.
**Action:** Profile attention memory scaling empirically:
1. Vary sequence length N (simulate different resolutions)
2. Track `mx.metal.get_peak_memory()` after each attention call
3. Plot peak memory vs N — quadratic = materialization, linear = streaming
**Impact:** If O(N^2) path is triggered at our token counts (~16K at 2048x2048), we may need to ensure contiguous 4D tensors entering SDPA — similar to Raiden129's patch but for Metal.
**Key finding from research:** Mask dtype must match Q/K/V dtype or entire computation upcasts to float32, halving bandwidth. Verify our attention masks (if any) match bf16 precision.
### 2. GPU-Side Preprocessing
**Problem:** Our engine likely does ImageNet normalization and resizing in numpy before converting to MLX arrays. On unified memory this is less costly than PCIe, but still leaves GPU ALUs idle during preprocessing.
**Action:**
- Audit `engine.py` preprocessing path
- Move normalize/resize to MLX operations
- Wrap in `@mx.compile` for kernel fusion (requires static input shapes)
**Caveat:** MLX lacks built-in bicubic resize. Manual implementation needed via `mx.meshgrid` + bilinear interpolation. Must use `mx.minimum`/`mx.maximum` for bounds clamping (no dynamic boolean indexing in MLX).
**Impact:** Moderate — eliminates numpy->mlx conversion overhead and uses GPU for parallel pixel operations.
### 3. Tiled Refiner: GPU Tensor Accumulators
**Problem:** Our tiled inference uses numpy accumulators for blend-weight averaging. This forces GPU->CPU transfer per tile and CPU->GPU for final result.
**Action:** Replace numpy accumulator arrays with MLX arrays. Accumulate blend weights and tile outputs entirely on GPU. Only convert final result to numpy at output.
**Impact:** Eliminates per-tile roundtrip. Especially significant at high tile counts (large images).
### 4. Non-Tiled Path: Strategic Memory Management
**Problem:** We do `gc.collect()` + `mx.metal.clear_cache()` in the tiling loop, but the non-tiled inference path doesn't have stage-boundary memory management.
**Action:** Add three-step cleanup protocol between pipeline stages in non-tiled forward:
1. `mx.eval()` on stage output tensors (force computation)
2. `del` intermediate tensors from previous stage
3. `gc.collect()` + `mx.metal.clear_cache()`
Insert at: encoder->decoder boundary, decoder->refiner boundary.
**Research confirms:** Just calling `mx.eval()` is insufficient. MLX's caching allocator hoards freed Metal buffers. Must explicitly `clear_cache()` to return memory to OS. This strictly bounds peak memory to the largest single stage rather than accumulated total.
**Impact:** Could significantly reduce peak memory in non-tiled path. Especially important for large img_size.
### 5. Token Routing (Experimental, Deferred)
**Problem:** Stage 2 has 16 blocks of global attention — dominates backbone compute. Many tokens are "easy" (solid FG/BG per alpha hint) and don't need full O(N^2) attention.
**Raiden129's approach:** LTRM module (LayerNorm->Linear->GELU->DWConv->Linear->ECA) at O(N) cost. Route by thresholding downsampled alpha hint. Zero-init fc2 weights for checkpoint compatibility.
**MLX challenge:** Three sparse processing strategies, each with tradeoffs:
| Strategy | Pros | Cons | Best When |
|---|---|---|---|
| `mx.where` + padding | Static shapes, compile-friendly | Doesn't reduce actual FLOPs | Low sparsity, many layers |
| Scatter + overflow bin | GPU-resident, physical reduction | Complex index arithmetic | Fixed-ratio routing |
| NumPy boolean indexing | Dynamic shapes, simple | CPU sync breaks async pipeline | Early, high-ratio culling (>60%) |
**Recommendation:** For CorridorKey, alpha hints typically have ~60-70% easy tokens. But routing happens at every block (16x in stage 2), so per-block CPU sync would be devastating. Best approach: `mx.where` with attention masking — keeps shapes static, compile-friendly, but note SDPA still processes padded tokens. Net benefit uncertain without benchmarking.
**Decision:** Defer until other optimizations landed. Needs careful profiling to verify compute savings > overhead.
### 6. Quantization (Future)
**Research findings:**
- **Int8:** ~50% memory reduction, ~1.8x speedup, <1% accuracy drop. Safe for all linear layers.
- **Int4:** ~75% memory reduction, ~2.4x speedup, 2-5% accuracy drop. Risky for matting precision.
- **MXFP4:** Best of both — 75% reduction, >2.4x speedup, <1.5% drop. M3/M4 optimized.
**Critical:** Leave Conv2d patch projection, LayerNorm, and softmax in full precision. These are pathologically sensitive to quantization in ViTs.
**For CorridorKey specifically:** Matting requires sub-pixel alpha precision. Int8 likely safe, Int4/MXFP4 needs quality validation on real footage. `mlx.nn.quantize` targets Linear/Embedding by default — Conv2d naturally excluded.
**Decision:** Defer. Profile Int8 as first candidate after other optimizations land.
## Key Decisions
1. **Attention profiling first** — verify O(N) vs O(N^2) behavior before optimizing attention path
2. **GPU preprocessing** — move normalize/resize to MLX, wrap in `@mx.compile`
3. **GPU tensor accumulators** — replace numpy accums in tiling with MLX arrays
4. **Non-tiled cache clearing** — add stage-boundary memory management protocol
5. **Token routing deferred** — MLX sparse processing constraints make ROI uncertain
6. **Quantization deferred** — Int8 first candidate, needs quality validation
## Priority Order
1. Non-tiled cache clearing (low effort, high impact on peak memory)
2. GPU tensor accumulators in tiling (medium effort, eliminates roundtrips)
3. Attention memory profiling (research, informs future work)
4. GPU preprocessing (medium effort, moderate speedup)
5. Token routing (high effort, uncertain ROI on MLX)
6. Int8 quantization (medium effort, needs quality validation)
## Open Questions
- What's our actual attention kernel path — O(N) or O(N^2) at 2048 resolution?
- How much time is spent in numpy preprocessing vs inference?
- Does `mx.metal.clear_cache()` between stages actually reduce peak memory in practice, or does lazy graph already handle it?
- For token routing: what % of tokens are "easy" in typical green screen footage?
- Int8 quantization: acceptable alpha precision loss threshold?

331
docs/plans/2026-03-09-feat-mlx-optimization-wave2-plan.md Normal file
View File

@ -0,0 +1,331 @@
---
title: "feat: MLX optimization wave 2 — SDPA, stage GC, slim forward"
type: feat
date: 2026-03-09
---
# MLX Optimization Wave 2
Second round of optimizations informed by PyTorch CorridorKey forks ([PR #104](https://github.com/nikopueringer/CorridorKey/pull/104), [Raiden129/CorridorKey_Test](https://github.com/Raiden129/CorridorKey_Test)) and [deep research on MLX ViT optimization](../MLX%20Vision%20Transformer%20Optimization%20Techniques.md).
**Brainstorm:** `docs/brainstorms/2026-03-09-pytorch-optimization-learnings-brainstorm.md`
**Wave 1 plan:** `docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md` (bf16, fused decode, tiled GC — all shipped)
## Overview
Six targeted optimizations, ordered by effort/impact ratio. Wave 1 achieved 12x peak memory reduction via tiled inference + deterministic GC. Wave 2 focuses on: fixing cleanup gaps, reducing computation graph size, and migrating to optimized kernels.
## Technical Approach
### Phase 1: Engine Cleanup Fixes (trivial, 15 min)
**Item 6: Add `mx.clear_cache()` to engine cleanup**
`engine.py:210` does `gc.collect()` but never calls `mx.clear_cache()`. The tiling module already does this (`tiling.py:164`), showing the pattern is known but wasn't applied to the non-tiled cleanup.
**Critical placement detail:** `mx.clear_cache()` must go **after postprocessing**, not at line 210. At line 210, `alpha_out`/`fg_out` are still live MLX arrays used by `postprocess_alpha/foreground` (which calls `np.array()` to transfer to CPU). Clearing cache while those references exist is pointless.
```python
# engine.py — after postprocessing extracts numpy arrays (~line 225)
del alpha_out, fg_out
gc.collect()
mx.clear_cache() # NEW: release Metal buffer cache
```
**Acceptance criteria:**
- [x] `mx.clear_cache()` called after all MLX arrays are consumed by postprocessing
- [x] `del` references to MLX arrays before cache clear
- [x] All existing tests pass (`uv run pytest`)
**Files:** `engine.py`
---
### Phase 2: Slim Forward Mode (low effort, 30 min)
**Item 3: Skip returning unused intermediate tensors**
`GreenFormer.forward()` (`corridorkey.py:103-113`) returns 9 dict keys. The engine uses at most 4 (`alpha_coarse`, `fg_coarse`, `alpha_final`, `fg_final`). The unused 5 (`alpha_logits`, `fg_logits`, `alpha_logits_up`, `fg_logits_up`, `delta_logits`) hold references that prevent MLX from freeing underlying buffers.
**Important clarification:** This is about **reference lifetime**, not computation skipping. All intermediates must still be computed to produce the final outputs. The savings come from not keeping references in the returned dict, allowing MLX to reclaim buffers sooner.
**Design decision:** Add `slim: bool = False` parameter to `GreenFormer.__init__`, not `__call__`. This avoids mx.compile recompilation from a runtime boolean changing the output dict shape. The `pipeline.infer()` API always returns the full dict for debugging. Only the engine sets `slim=True`.
```python
# corridorkey.py
class GreenFormer(nn.Module):
def __init__(self, ..., slim: bool = False):
self.slim = slim
...
def __call__(self, x):
...
if self.slim:
return {
"alpha_coarse": alpha_coarse,
"fg_coarse": fg_coarse,
"alpha_final": alpha_final,
"fg_final": fg_final,
}
return { # full dict for debugging/testing
"alpha_logits": ...,
...
}
```
**Acceptance criteria:**
- [x] `slim=True` returns 4-key dict, `slim=False` returns full 9-key dict
- [x] Engine passes `slim=True` via `load_model()`
- [x] `pipeline.infer()` keeps `slim=False` (preserves debugging API)
- [x] Parity tests run with `slim=False` (validate intermediate keys unchanged)
- [x] Add test: slim output matches corresponding keys from full output
**Files:** `corridorkey.py`, `engine.py`, `pipeline.py`, new test in `tests/`
---
### Phase 3: Non-Tiled Stage-Boundary Memory Management (medium effort, 1 hr)
**Item 1: Add intermediate graph materialization between encoder/decoder/refiner**
The non-tiled path builds one massive computation graph across all 24 Hiera blocks + 2 decoders + refiner before any materialization. Inserting `mx.eval()` at stage boundaries lets MLX free intermediate graph nodes.
**mx.compile interaction (CRITICAL):**
`mx.compile` wraps the entire `GreenFormer.__call__`. Inserting `mx.eval()` inside a compiled function breaks the compilation graph. Solution: add a `_compiled: bool` flag set by `load_model()` when compile is enabled. Skip stage-boundary GC when compiled.
```python
# corridorkey.py — inside __call__
features = self.backbone(x)
if not self._compiled:
mx.eval(features) # materialize backbone output
gc.collect()
mx.clear_cache() # free backbone intermediate graph
# ... decoders ...
coarse = {...}
if not self._compiled:
mx.eval(coarse)
gc.collect()
mx.clear_cache()
# ... refiner ...
```
**Key question answered:** Is this worth it for full-frame? The benchmark shows 27.6 GB peak for full-frame. The backbone's intermediate computation graph (24 blocks of attention + MLP, accumulated lazily) likely dominates. Breaking it into 3 smaller graphs (backbone, decoders, refiner) should reduce peak significantly by allowing MLX to reclaim backbone intermediates before running decoders.
**Acceptance criteria:**
- [x] `_compiled` flag set by `compile_model()` when compile=True
- [x] Stage-boundary GC only runs when `_compiled=False`
- [ ] Full-frame peak memory measurably decreases (measure with `mx.metal.get_peak_memory()`)
- [x] Compiled path behavior unchanged (no mx.eval calls inside graph)
- [x] All parity tests pass
**Files:** `corridorkey.py`, `pipeline.py`
---
### Phase 4: SDPA Migration (medium effort, 1-2 hr)
**Item 2: Replace manual attention with `mx.fast.scaled_dot_product_attention`**
Current (`hiera.py:288-290`):
```python
attn = (q * self.scale) @ mx.transpose(k, ...)
attn = mx.softmax(attn, axis=-1)
x = attn @ v
```
This materializes the full `(B, heads, windows, tokens, tokens)` attention matrix. `mx.fast.scaled_dot_product_attention` fuses this into a single Metal kernel.
**Pre-requisites (spike required, 15 min):**
Before implementing, verify with a quick spike:
1. **5D input handling:** Current Q/K/V are 5D `(B, heads, num_windows, tokens_per_window, head_dim)`. SDPA expects 4D `(batch, heads, seq_len, head_dim)`. Must fold `num_windows` into batch dim: reshape to `(B * num_windows, heads, tokens_per_window, head_dim)`, call SDPA, reshape back.
2. **Asymmetric Q/K/V lengths:** At stride blocks (indices 2, 5, 21), Q is max-pooled to fewer tokens than K/V. Verify `mx.fast.scaled_dot_product_attention` supports `q_len != kv_len` (standard for cross-attention, should work).
3. **Scale factor application:** Current code does `(q * scale) @ k.T`. SDPA does `(q @ k.T) * scale` internally. Mathematically equivalent in fp32 but different rounding in bf16. Measure parity impact.
**Implementation:**
```python
# hiera.py — MaskUnitAttention.__call__
# Fold windows into batch for SDPA
B, H, W, T, D = q.shape
q_4d = q.reshape(B * W, H, -1, D) # (B*windows, heads, tokens, head_dim)
k_4d = k.reshape(B * W, H, T, D)
v_4d = v.reshape(B * W, H, T, D)
x = mx.fast.scaled_dot_product_attention(
q_4d, k_4d, v_4d, scale=self.scale
)
x = x.reshape(B, H, W, -1, D) # unfold windows
```
**Attention matrix sizes (reference):**
| Stage | Blocks | mask_unit_attn | Tokens per window | Attention matrix |
|---|---|---|---|---|
| 0 | 0-1 | True | 64 | 64x64 |
| 1 | 2-4 | True | 16 | 16x16 |
| 2 | 5-20 | False | 256 | 256x256 |
| 3 | 21-23 | False | 64 | 64x64 |
Sizes are modest thanks to Hiera's unrolling. Memory savings per SDPA call are small (~4MB for 256x256). Primary benefit is **kernel fusion** — one Metal dispatch vs three (matmul + softmax + matmul).
**Acceptance criteria:**
- [x] Spike confirms SDPA supports: 4D input, asymmetric Q/K, bf16
- [x] All attention calls use `mx.fast.scaled_dot_product_attention`
- [x] Window dim folded into batch before SDPA, unfolded after
- [x] Stride blocks (Q pooling) produce correct output with asymmetric lengths
- [x] E2E parity within `1e-3` tolerance (existing `PARITY_TOL_E2E`)
- [x] Backbone stage parity regressions documented if any
**Files:** `hiera.py`
---
### Phase 5: GPU Preprocessing (medium effort, 1-2 hr)
**Item 5: Move ImageNet normalize + resize from numpy/PIL to MLX**
Current flow (`io/image.py:48-68`, `engine.py:153-181`):
1. `image.astype(np.float32) / 255.0` — numpy
2. PIL bicubic resize — numpy/PIL
3. ImageNet normalize + concat — numpy
4. Single `mx.array()` call — boundary
**Scope limitation:** Full-frame path only. For tiled inference, the full-res image must stay accessible for tile slicing. Moving full-res preprocessing to GPU would allocate the entire image on GPU before tiling begins — counterproductive.
**Implementation approach:**
```python
# New: io/preprocess_mlx.py
def preprocess_mlx(image_uint8: np.ndarray, mask: np.ndarray, img_size: int) -> mx.array:
"""GPU-side preprocessing for full-frame inference."""
# Convert to MLX early
img = mx.array(image_uint8).astype(mx.float32) / 255.0 # (H, W, 3)
mask = mx.array(mask).astype(mx.float32) # (H, W, 1)
# Resize using nn.Upsample (bilinear, not bicubic)
# Note: scale_factor = img_size / current_size
img = mx.expand_dims(img, axis=0) # (1, H, W, 3) NHWC
img = upsample(img, target_size=img_size)
mask = mx.expand_dims(mask, axis=0)
mask = upsample(mask, target_size=img_size)
# ImageNet normalize
mean = mx.array([0.485, 0.456, 0.406]) # (3,)
std = mx.array([0.229, 0.224, 0.225])
img = (img - mean) / std
# Concat
return mx.concatenate([img, mask], axis=-1) # (1, H, W, 4) NHWC
```
**Caveat — resize quality:** PIL `BICUBIC` and MLX `nn.Upsample(mode="cubic")` use different cubic kernels. May introduce parity differences. Run a measurement spike: compare PIL bicubic vs MLX cubic on the golden test image, measure max absolute difference.
**Acceptance criteria:**
- [x] Full-frame normalize+concat runs on GPU (resize stays CPU/PIL)
- [x] Tiled path unchanged (keeps numpy preprocessing)
- [x] Resize quality spike: PIL bicubic vs MLX cubic max diff = 0.22 (upscale), 0.59 (downscale)
- [x] Resize diff >> 1e-3: kept PIL resize, moved only normalize+concat to GPU
- [x] E2E parity within tolerance (preprocess_mlx vs preprocess: 0.0 diff)
- [ ] Optional: wrap in `@mx.compile` for static input sizes
**Files:** New `io/preprocess_mlx.py`, `engine.py`
---
### Phase 6: GPU Tensor Accumulators in Tiling (nice-to-have, deferred)
**Item 4: Replace numpy accumulators with MLX arrays**
Current (`tiling.py:120-122`): Three numpy float32 arrays at full resolution. Per-tile results are transferred via `np.array(out["alpha_final"][0])`.
**Why deferred:**
1. MLX lacks in-place slice assignment (`arr[y:y_end, x:x_end] += tile`)
2. `mx.scatter_add` or equivalent needs API verification
3. GPU->CPU transfer per tile is ~1MB at 512x512 — microseconds on unified memory
4. Wave 1 plan already flagged this as "best-effort, 30min timebox" with numpy fallback
5. Total transfer for 16 tiles (2048x2048) is ~16MB — negligible vs compute time
**If pursued later:** Investigate `mx.scatter_add` API, or accumulate using full-size zero tensors with padded tiles (wasteful but avoids slice assignment).
---
## Alternative Approaches Considered
### Token routing (from Raiden129)
Route "easy" tokens to lightweight LTRM module, skip expensive attention. **Rejected for now:** MLX lacks dynamic boolean indexing. `mx.where` + padding doesn't reduce FLOPs. NumPy CPU fallback per-block (16x in stage 2) would destroy async pipeline. Net benefit uncertain. Revisit if profiling shows attention is the bottleneck.
### Int8/MXFP4 quantization
50-75% weight memory reduction. **Deferred:** Matting requires sub-pixel alpha precision. Int8 likely safe (<1% drop) but needs quality validation on real footage. Leave Conv2d, LayerNorm, softmax in full precision. Profile as separate future work.
### Full weight-level bf16 casting
Cast all model weights to bf16 at load time (halves parameter memory). **Not yet explored:** Currently only activations use bf16. Would need to verify backbone parity in bf16 weights (16 blocks = most drift risk). Lower priority than computation graph optimizations.
## Acceptance Criteria
### Functional Requirements
- [ ] All 94 existing tests pass
- [ ] E2E parity within `1e-3` tolerance
- [ ] No regression in tiled inference behavior
- [ ] Engine cleanup releases all Metal buffers after postprocessing
### Non-Functional Requirements
- [ ] Full-frame peak memory measurably decreases with stage-boundary GC (measure with `mx.metal.get_peak_memory()`)
- [ ] No latency regression > 5% in any mode
- [ ] `uv run ruff check .` passes
- [ ] `uv run ruff format --check .` passes
- [ ] `uv run ty check` passes
### Quality Gates
- [ ] SDPA spike completed before Phase 4 implementation
- [ ] Resize quality spike completed before Phase 5 implementation
- [ ] Each phase merged separately with passing CI
## Implementation Order & Dependencies
```
Phase 1 (engine cleanup) ─── no deps ──────────────────── 15 min
Phase 2 (slim forward) ─── no deps ──────────────────── 30 min
Phase 3 (stage GC) ─── depends on Phase 2 (slim changes what's in dict) ── 1 hr
Phase 4 (SDPA) ─── no deps (backbone only) ──── spike 15m + impl 1-2 hr
Phase 5 (GPU preprocess) ── no deps (I/O only) ────────── spike 15m + impl 1-2 hr
Phase 6 (GPU accum) ─── deferred ─────────────────── (skip)
```
Phases 1-2 and Phase 4 can run in parallel. Phase 3 should follow Phase 2.
## References
### Internal
- Wave 1 plan: `docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md`
- Brainstorm: `docs/brainstorms/2026-03-09-pytorch-optimization-learnings-brainstorm.md`
- Deep research: `docs/MLX Vision Transformer Optimization Techniques.md`
- Attention impl: `src/corridorkey_mlx/model/hiera.py:267-297`
- Forward pass: `src/corridorkey_mlx/model/corridorkey.py:56-113`
- Engine cleanup: `src/corridorkey_mlx/engine.py:187-210`
- Tiling accumulators: `src/corridorkey_mlx/inference/tiling.py:120-171`
- Preprocessing: `src/corridorkey_mlx/io/image.py:48-68`
- Pipeline defaults: `src/corridorkey_mlx/inference/pipeline.py:31-58`
### External
- PR #104: https://github.com/nikopueringer/CorridorKey/pull/104
- Raiden129 fork: https://github.com/Raiden129/CorridorKey_Test
- MLX SDPA docs: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.fast.scaled_dot_product_attention.html
- MLX FlashAttention issue: https://github.com/ml-explore/mlx/issues/129
- MLX memory management: https://github.com/ml-explore/mlx/issues/742
## Open Questions
- Stage-boundary GC: measurable peak memory reduction in full-frame, or negligible?
- SDPA spike: does MLX SDPA handle 5D->4D reshape + asymmetric Q/K correctly?
- Resize quality: PIL bicubic vs MLX cubic — within parity tolerance?
- Weight-level bf16: safe for Hiera backbone (16 blocks of drift)?
- Int8 quantization: acceptable alpha precision threshold for matting?