feat: add optimization toggle flags + benchmark matrix script

use_sdpa, stage_gc flags propagated through backbone/model/pipeline.
bench_optimizations.py runs exhaustive/ablation/key sweep of all 6
toggles (slim, stage_gc, sdpa, bf16, fused_decode, gpu_preprocess)
with latency + peak memory reporting.

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

scripts/bench_optimizations.py

@@ -0,0 +1,450 @@
+#!/usr/bin/env python3
+"""Benchmark matrix for wave 2 optimizations.
+
+Runs an exhaustive matrix of optimization toggles and reports latency + peak
+memory for each configuration. Helps isolate the impact of each optimization.
+
+Usage:
+    uv run python scripts/bench_optimizations.py
+    uv run python scripts/bench_optimizations.py --resolution 512 --bench-runs 5
+    uv run python scripts/bench_optimizations.py --checkpoint ckpt.safetensors
+
+All mx.eval() calls are MLX array materialization, NOT Python eval().
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import itertools
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import mlx.core as mx
+from rich.console import Console
+from rich.table import Table
+
+# Ensure package is importable when running from repo root
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from corridorkey_mlx.io.image import preprocess as preprocess_numpy
+from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx
+from corridorkey_mlx.model.corridorkey import GreenFormer
+
+console = Console()
+
+# ---------------------------------------------------------------------------
+# Optimization flags
+# ---------------------------------------------------------------------------
+
+TOGGLE_NAMES = ["slim", "stage_gc", "sdpa", "bf16", "fused_decode", "gpu_preprocess"]
+
+
+@dataclass
+class OptConfig:
+    """Single optimization configuration to benchmark."""
+
+    slim: bool = True
+    stage_gc: bool = True
+    sdpa: bool = True
+    bf16: bool = True
+    fused_decode: bool = True
+    gpu_preprocess: bool = True
+
+    @property
+    def label(self) -> str:
+        flags = []
+        for name in TOGGLE_NAMES:
+            val = getattr(self, name)
+            flags.append(f"{name}={'on' if val else 'off'}")
+        return " | ".join(flags)
+
+    @property
+    def short_label(self) -> str:
+        parts = []
+        for name in TOGGLE_NAMES:
+            if getattr(self, name):
+                parts.append(name)
+        return "+".join(parts) if parts else "baseline"
+
+
+@dataclass
+class BenchResult:
+    """Result from a single benchmark configuration."""
+
+    config: OptConfig
+    resolution: int
+    warmup_ms: float = 0.0
+    median_ms: float = 0.0
+    min_ms: float = 0.0
+    peak_memory_mb: float = 0.0
+    all_times: list[float] = field(default_factory=list)
+    error: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# Benchmark logic
+# ---------------------------------------------------------------------------
+
+
+def build_model(
+    config: OptConfig,
+    img_size: int,
+    checkpoint: Path | None,
+) -> GreenFormer:
+    """Build model with the given optimization config."""
+    dtype = mx.bfloat16 if config.bf16 else mx.float32
+    model = GreenFormer(
+        img_size=img_size,
+        dtype=dtype,
+        fused_decode=config.fused_decode,
+        slim=config.slim,
+        use_sdpa=config.sdpa,
+        stage_gc=config.stage_gc,
+    )
+    if checkpoint and checkpoint.exists():
+        model.load_checkpoint(checkpoint)
+    else:
+        model.eval()
+        mx.eval(model.parameters())  # noqa: S307
+    return model
+
+
+def make_input(img_size: int, gpu_preprocess: bool) -> mx.array:
+    """Create input tensor, optionally via GPU preprocessing path."""
+    import numpy as np
+
+    np.random.seed(42)
+    rgb_f32 = np.random.rand(img_size, img_size, 3).astype(np.float32)
+    mask_f32 = np.random.rand(img_size, img_size, 1).astype(np.float32)
+
+    if gpu_preprocess:
+        return preprocess_mlx(rgb_f32, mask_f32)
+    return preprocess_numpy(rgb_f32, mask_f32)
+
+
+def time_inference(
+    model: GreenFormer,
+    x: mx.array,
+    warmup_runs: int,
+    bench_runs: int,
+) -> tuple[float, float, float, list[float]]:
+    """Run warmup + benchmark, return (warmup_ms, median_ms, min_ms, all_times)."""
+
+    def run() -> None:
+        out = model(x)
+        mx.eval(out)  # noqa: S307
+
+    # Warmup
+    start = time.perf_counter()
+    run()
+    warmup_ms = (time.perf_counter() - start) * 1000.0
+
+    for _ in range(warmup_runs - 1):
+        run()
+
+    # Benchmark
+    times: list[float] = []
+    for _ in range(bench_runs):
+        start = time.perf_counter()
+        run()
+        times.append((time.perf_counter() - start) * 1000.0)
+
+    times.sort()
+    median_ms = times[len(times) // 2]
+    min_ms = times[0]
+    return warmup_ms, median_ms, min_ms, times
+
+
+def measure_peak_memory(model: GreenFormer, x: mx.array) -> float:
+    """Measure peak Metal memory for a single forward pass (MB)."""
+    gc.collect()
+    mx.clear_cache()
+
+    # Use non-deprecated API if available, fall back to mx.metal.*
+    reset_fn = getattr(mx, "reset_peak_memory", None) or getattr(
+        mx.metal, "reset_peak_memory", None
+    )
+    get_fn = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
+    if not reset_fn or not get_fn:
+        return 0.0
+
+    reset_fn()
+
+    out = model(x)
+    mx.eval(out)  # noqa: S307
+
+    peak_bytes = get_fn()
+
+    del out
+    gc.collect()
+    mx.clear_cache()
+
+    return peak_bytes / (1024 * 1024)
+
+
+def bench_config(
+    config: OptConfig,
+    resolution: int,
+    checkpoint: Path | None,
+    warmup_runs: int,
+    bench_runs: int,
+) -> BenchResult:
+    """Benchmark a single optimization configuration."""
+    result = BenchResult(config=config, resolution=resolution)
+
+    try:
+        model = build_model(config, resolution, checkpoint)
+        x = make_input(resolution, config.gpu_preprocess)
+        mx.eval(x)  # noqa: S307
+
+        warmup_ms, median_ms, min_ms, all_times = time_inference(model, x, warmup_runs, bench_runs)
+        result.warmup_ms = warmup_ms
+        result.median_ms = median_ms
+        result.min_ms = min_ms
+        result.all_times = all_times
+
+        # Measure peak memory with a fresh model to avoid cached state
+        model_fresh = build_model(config, resolution, checkpoint)
+        result.peak_memory_mb = measure_peak_memory(model_fresh, x)
+
+        # Cleanup
+        del model, model_fresh, x
+        gc.collect()
+        mx.clear_cache()
+
+    except Exception as e:
+        result.error = str(e)
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Matrix generation
+# ---------------------------------------------------------------------------
+
+
+def generate_matrix(sweep: str) -> list[OptConfig]:
+    """Generate benchmark configurations.
+
+    Args:
+        sweep: "full" for exhaustive 2^6 matrix, "ablation" for toggle-one-off,
+               "key" for important combinations only.
+    """
+    all_on = OptConfig()
+
+    if sweep == "ablation":
+        # All-on baseline + toggle each flag off one at a time
+        configs = [all_on]
+        for name in TOGGLE_NAMES:
+            off = OptConfig(**{**vars(all_on), name: False})
+            configs.append(off)
+        # Also add all-off baseline
+        configs.append(
+            OptConfig(
+                slim=False,
+                stage_gc=False,
+                sdpa=False,
+                bf16=False,
+                fused_decode=False,
+                gpu_preprocess=False,
+            )
+        )
+        return configs
+
+    if sweep == "full":
+        # Exhaustive 2^6 = 64 configurations
+        configs = []
+        for combo in itertools.product([True, False], repeat=len(TOGGLE_NAMES)):
+            kwargs = dict(zip(TOGGLE_NAMES, combo, strict=True))
+            configs.append(OptConfig(**kwargs))
+        return configs
+
+    # "key" — important configurations
+    return [
+        # All off (wave 1 baseline)
+        OptConfig(
+            slim=False,
+            stage_gc=False,
+            sdpa=False,
+            bf16=False,
+            fused_decode=False,
+            gpu_preprocess=False,
+        ),
+        # Wave 1 only (bf16 + fused_decode)
+        OptConfig(
+            slim=False,
+            stage_gc=False,
+            sdpa=False,
+            bf16=True,
+            fused_decode=True,
+            gpu_preprocess=False,
+        ),
+        # Wave 2 only (slim + stage_gc + sdpa + gpu_preprocess)
+        OptConfig(
+            slim=True,
+            stage_gc=True,
+            sdpa=True,
+            bf16=False,
+            fused_decode=False,
+            gpu_preprocess=True,
+        ),
+        # All on
+        all_on,
+        # SDPA only
+        OptConfig(
+            slim=False,
+            stage_gc=False,
+            sdpa=True,
+            bf16=False,
+            fused_decode=False,
+            gpu_preprocess=False,
+        ),
+        # Stage GC only
+        OptConfig(
+            slim=False,
+            stage_gc=True,
+            sdpa=False,
+            bf16=False,
+            fused_decode=False,
+            gpu_preprocess=False,
+        ),
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Output
+# ---------------------------------------------------------------------------
+
+
+def print_results(results: list[BenchResult]) -> None:
+    """Print results as a rich table."""
+    table = Table(title="Optimization Benchmark Matrix")
+
+    table.add_column("Config", max_width=50)
+    table.add_column("Res", justify="right")
+    table.add_column("Warmup (ms)", justify="right")
+    table.add_column("Median (ms)", justify="right")
+    table.add_column("Min (ms)", justify="right")
+    table.add_column("Peak Mem (MB)", justify="right")
+    table.add_column("Status", justify="center")
+
+    # Sort by resolution then median time
+    results.sort(key=lambda r: (r.resolution, r.median_ms))
+
+    for r in results:
+        if r.error:
+            table.add_row(
+                r.config.short_label,
+                str(r.resolution),
+                "-",
+                "-",
+                "-",
+                "-",
+                f"[red]FAIL: {r.error[:30]}[/red]",
+            )
+        else:
+            table.add_row(
+                r.config.short_label,
+                str(r.resolution),
+                f"{r.warmup_ms:.1f}",
+                f"{r.median_ms:.1f}",
+                f"{r.min_ms:.1f}",
+                f"{r.peak_memory_mb:.0f}" if r.peak_memory_mb > 0 else "N/A",
+                "[green]OK[/green]",
+            )
+
+    console.print(table)
+
+    # Print speedup summary vs baseline
+    baselines = {r.resolution: r for r in results if r.config.short_label == "baseline"}
+    if baselines:
+        console.print("\n[bold]Speedup vs baseline (all-off):[/bold]")
+        for r in results:
+            if r.error or r.config.short_label == "baseline":
+                continue
+            baseline = baselines.get(r.resolution)
+            if baseline and baseline.median_ms > 0:
+                speedup = baseline.median_ms / r.median_ms
+                mem_ratio = (
+                    f"{r.peak_memory_mb / baseline.peak_memory_mb:.2f}x"
+                    if baseline.peak_memory_mb > 0 and r.peak_memory_mb > 0
+                    else "N/A"
+                )
+                console.print(
+                    f"  {r.config.short_label:40s} "
+                    f"{speedup:.2f}x speed, {mem_ratio} memory "
+                    f"({r.resolution})"
+                )
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark matrix for wave 2 optimizations")
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=Path("checkpoints/corridorkey_mlx.safetensors"),
+    )
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        nargs="+",
+        default=[512],
+        help="Resolution(s) to benchmark",
+    )
+    parser.add_argument(
+        "--sweep",
+        choices=["key", "ablation", "full"],
+        default="ablation",
+        help="key=6 configs, ablation=8 configs (toggle one off), full=64 configs",
+    )
+    parser.add_argument("--warmup-runs", type=int, default=3)
+    parser.add_argument("--bench-runs", type=int, default=10)
+    args = parser.parse_args()
+
+    ckpt = args.checkpoint if args.checkpoint.exists() else None
+    configs = generate_matrix(args.sweep)
+
+    console.print("[bold]CorridorKey MLX — Optimization Benchmark Matrix[/bold]")
+    console.print(f"  Weights: {args.checkpoint if ckpt else 'random (no checkpoint)'}")
+    console.print(f"  Resolutions: {args.resolution}")
+    console.print(f"  Sweep: {args.sweep} ({len(configs)} configs)")
+    console.print(f"  Warmup: {args.warmup_runs}, Bench: {args.bench_runs}")
+    console.print()
+
+    all_results: list[BenchResult] = []
+    total = len(configs) * len(args.resolution)
+    idx = 0
+
+    for res in args.resolution:
+        for config in configs:
+            idx += 1
+            console.print(
+                f"  [{idx}/{total}] {res}x{res} — {config.short_label}...",
+                end=" ",
+            )
+            result = bench_config(config, res, ckpt, args.warmup_runs, args.bench_runs)
+            if result.error:
+                console.print("[red]FAIL[/red]")
+            else:
+                console.print(
+                    f"[green]{result.median_ms:.1f}ms[/green] "
+                    f"(peak: {result.peak_memory_mb:.0f}MB)"
+                    if result.peak_memory_mb > 0
+                    else f"[green]{result.median_ms:.1f}ms[/green]"
+                )
+            all_results.append(result)
+
+    console.print()
+    print_results(all_results)
+
+
+if __name__ == "__main__":
+    main()

src/corridorkey_mlx/inference/pipeline.py

@@ -36,6 +36,8 @@ def load_model(
     dtype: mx.Dtype = mx.bfloat16,
     fused_decode: bool = True,
     slim: bool = False,
+    use_sdpa: bool = True,
+    stage_gc: bool = True,
 ) -> GreenFormer:
     """Build GreenFormer and load weights from safetensors checkpoint.
 
@@ -53,8 +55,17 @@ def load_model(
             Metal dispatch calls. Bit-exact with unfused path.
         slim: If True, forward returns only 4 final keys (drops intermediates).
             Reduces reference lifetime so MLX can reclaim buffers sooner.
+        use_sdpa: If True, use mx.fast.scaled_dot_product_attention in backbone.
+        stage_gc: If True, materialize + GC at backbone/decoder/refiner boundaries.
     """
-    model = GreenFormer(img_size=img_size, dtype=dtype, fused_decode=fused_decode, slim=slim)
+    model = GreenFormer(
+        img_size=img_size,
+        dtype=dtype,
+        fused_decode=fused_decode,
+        slim=slim,
+        use_sdpa=use_sdpa,
+        stage_gc=stage_gc,
+    )
     model.load_checkpoint(checkpoint)
     if compile:
         model = compile_model(model, shapeless=shapeless)

src/corridorkey_mlx/model/corridorkey.py

@@ -38,13 +38,16 @@ class GreenFormer(nn.Module):
         dtype: mx.Dtype = mx.float32,
         fused_decode: bool = False,
         slim: bool = False,
+        use_sdpa: bool = True,
+        stage_gc: bool = True,
     ) -> None:
         super().__init__()
         self._compute_dtype = dtype
         self._fused_decode = fused_decode
         self._slim = slim
+        self._stage_gc = stage_gc
         self._compiled = False
-        self.backbone = HieraBackbone(img_size=img_size)
+        self.backbone = HieraBackbone(img_size=img_size, use_sdpa=use_sdpa)
         self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1)
         self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3)
         self.refiner = CNNRefinerModule()
@@ -73,7 +76,7 @@ class GreenFormer(nn.Module):
 
         # Materialize backbone output so MLX can free intermediate graph nodes.
         # NOTE: mx.eval is MLX array materialization, not Python eval()
-        if not self._compiled:
+        if self._stage_gc and not self._compiled:
             mx.eval(features)  # noqa: S307
             gc.collect()
             mx.clear_cache()
@@ -104,7 +107,7 @@ class GreenFormer(nn.Module):
 
         # Materialize decoder output so MLX can free decoder graph nodes.
         # NOTE: mx.eval is MLX array materialization, not Python eval()
-        if not self._compiled:
+        if self._stage_gc and not self._compiled:
             mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up)  # noqa: S307
             gc.collect()
             mx.clear_cache()

src/corridorkey_mlx/model/hiera.py

@@ -251,6 +251,7 @@ class MaskUnitAttention(nn.Module):
         q_stride: int = 1,
         window_size: int = 0,
         use_mask_unit_attn: bool = False,
+        use_sdpa: bool = True,
     ) -> None:
         super().__init__()
         self.dim_out = dim_out
@@ -260,6 +261,7 @@ class MaskUnitAttention(nn.Module):
         self.scale = self.head_dim**-0.5
         self.window_size = window_size
         self.use_mask_unit_attn = use_mask_unit_attn
+        self._use_sdpa = use_sdpa
 
         self.qkv = nn.Linear(dim, 3 * dim_out)
         self.proj = nn.Linear(dim_out, dim_out)
@@ -284,26 +286,32 @@ class MaskUnitAttention(nn.Module):
             q = q.reshape(batch_size, self.heads, num_windows, self.q_stride, -1, self.head_dim)
             q = mx.max(q, axis=3)
 
-        # Fold window dim into batch for SDPA (expects 4D).
+        if self._use_sdpa:
-        # Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows
+            # Fold window dim into batch for SDPA (expects 4D).
-        # next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D).
+            # Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows
-        q_tokens = q.shape[3]
+            # next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D).
-        kv_tokens = k.shape[3]
+            q_tokens = q.shape[3]
-        q_4d = mx.transpose(q, axes=(0, 2, 1, 3, 4)).reshape(
+            kv_tokens = k.shape[3]
-            batch_size * num_windows, self.heads, q_tokens, self.head_dim
+            q_4d = mx.transpose(q, axes=(0, 2, 1, 3, 4)).reshape(
-        )
+                batch_size * num_windows, self.heads, q_tokens, self.head_dim
-        k_4d = mx.transpose(k, axes=(0, 2, 1, 3, 4)).reshape(
+            )
-            batch_size * num_windows, self.heads, kv_tokens, self.head_dim
+            k_4d = mx.transpose(k, axes=(0, 2, 1, 3, 4)).reshape(
-        )
+                batch_size * num_windows, self.heads, kv_tokens, self.head_dim
-        v_4d = mx.transpose(v, axes=(0, 2, 1, 3, 4)).reshape(
+            )
-            batch_size * num_windows, self.heads, kv_tokens, self.head_dim
+            v_4d = mx.transpose(v, axes=(0, 2, 1, 3, 4)).reshape(
-        )
+                batch_size * num_windows, self.heads, kv_tokens, self.head_dim
+            )
 
-        x = mx.fast.scaled_dot_product_attention(q_4d, k_4d, v_4d, scale=self.scale)
+            x = mx.fast.scaled_dot_product_attention(q_4d, k_4d, v_4d, scale=self.scale)
 
-        # Unfold: (B*W, H, T, D) -> (B, W, H, T, D) -> (B, H, W, T, D)
+            # Unfold: (B*W, H, T, D) -> (B, W, H, T, D) -> (B, H, W, T, D)
-        x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim)
+            x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim)
-        x = mx.transpose(x, axes=(0, 2, 1, 3, 4))
+            x = mx.transpose(x, axes=(0, 2, 1, 3, 4))
+        else:
+            # Manual scaled dot-product attention (original implementation)
+            attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3))
+            attn = mx.softmax(attn, axis=-1)
+            x = attn @ v
 
         # [B, heads, num_windows, tokens, head_dim] -> [B, tokens, num_windows, heads, head_dim]
         # -> [B, N', dim_out]  (matches PyTorch transpose(1, 3))
@@ -328,6 +336,7 @@ class HieraBlock(nn.Module):
         q_stride: int = 1,
         window_size: int = 0,
         use_mask_unit_attn: bool = False,
+        use_sdpa: bool = True,
     ) -> None:
         super().__init__()
         self.dim = dim
@@ -340,7 +349,7 @@ class HieraBlock(nn.Module):
             self.proj = nn.Linear(dim, dim_out)
 
         self.attn = MaskUnitAttention(
-            dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn
+            dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn, use_sdpa=use_sdpa
         )
 
         self.norm2 = nn.LayerNorm(dim_out)
@@ -369,7 +378,7 @@ class HieraBackbone(nn.Module):
     Hardcoded for hiera_base_plus_224 config with 4-channel input.
     """
 
-    def __init__(self, img_size: int = 512) -> None:
+    def __init__(self, img_size: int = 512, use_sdpa: bool = True) -> None:
         super().__init__()
         self.img_size = img_size
 
@@ -434,6 +443,7 @@ class HieraBackbone(nn.Module):
                 q_stride=(flat_q_stride if i in q_pool_blocks else 1),
                 window_size=flat_mu_size,
                 use_mask_unit_attn=use_mu_attn,
+                use_sdpa=use_sdpa,
             )
             self.blocks.append(block)
             embed_dim = dim_out