diff --git a/scripts/bench_optimizations.py b/scripts/bench_optimizations.py new file mode 100644 index 0000000..017fda1 --- /dev/null +++ b/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() diff --git a/src/corridorkey_mlx/inference/pipeline.py b/src/corridorkey_mlx/inference/pipeline.py index be94813..37e7909 100644 --- a/src/corridorkey_mlx/inference/pipeline.py +++ b/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) diff --git a/src/corridorkey_mlx/model/corridorkey.py b/src/corridorkey_mlx/model/corridorkey.py index 979dda2..a09cc62 100644 --- a/src/corridorkey_mlx/model/corridorkey.py +++ b/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() diff --git a/src/corridorkey_mlx/model/hiera.py b/src/corridorkey_mlx/model/hiera.py index 46a5809..d48ca95 100644 --- a/src/corridorkey_mlx/model/hiera.py +++ b/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). - # Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows - # next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D). - q_tokens = q.shape[3] - kv_tokens = k.shape[3] - 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 - ) - v_4d = mx.transpose(v, axes=(0, 2, 1, 3, 4)).reshape( - batch_size * num_windows, self.heads, kv_tokens, self.head_dim - ) + if self._use_sdpa: + # Fold window dim into batch for SDPA (expects 4D). + # Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows + # next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D). + q_tokens = q.shape[3] + kv_tokens = k.shape[3] + 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 + ) + 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) - x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim) - x = mx.transpose(x, axes=(0, 2, 1, 3, 4)) + # 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 = 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