corridorkey-mrp-mlx/docs/plans/2026-03-09-feat-mlx-optimization-wave2-plan.md

title	type	date
feat: MLX optimization wave 2 — SDPA, stage GC, slim forward	feat	2026-03-09

MLX Optimization Wave 2

Second round of optimizations informed by PyTorch CorridorKey forks (PR #104, Raiden129/CorridorKey_Test) and deep research on MLX ViT optimization.

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.

# 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:

• mx.clear_cache() called after all MLX arrays are consumed by postprocessing
• del references to MLX arrays before cache clear
• 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.

# 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:

• slim=True returns 4-key dict, slim=False returns full 9-key dict
• Engine passes slim=True via load_model()
• pipeline.infer() keeps slim=False (preserves debugging API)
• Parity tests run with slim=False (validate intermediate keys unchanged)
• 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.

# 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:

• _compiled flag set by compile_model() when compile=True
• Stage-boundary GC only runs when _compiled=False
• Full-frame peak memory measurably decreases (measure with mx.metal.get_peak_memory())
• Compiled path behavior unchanged (no mx.eval calls inside graph)
• 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):

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:

# 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:

• Spike confirms SDPA supports: 4D input, asymmetric Q/K, bf16
• All attention calls use mx.fast.scaled_dot_product_attention
• Window dim folded into batch before SDPA, unfolded after
• Stride blocks (Q pooling) produce correct output with asymmetric lengths
• E2E parity within 1e-3 tolerance (existing PARITY_TOL_E2E)
• 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:

# 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:

• Full-frame normalize+concat runs on GPU (resize stays CPU/PIL)
• Tiled path unchanged (keeps numpy preprocessing)
• Resize quality spike: PIL bicubic vs MLX cubic max diff = 0.22 (upscale), 0.59 (downscale)
• Resize diff >> 1e-3: kept PIL resize, moved only normalize+concat to GPU
• 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?