feat: stage-boundary GC between backbone/decoder/refiner

mx.eval + gc + clear_cache at backbone→decoder and decoder→refiner
boundaries in eager mode. _compiled flag skips GC under mx.compile.
Compilation tests updated to use compile_model() API.

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

src/corridorkey_mlx/inference/pipeline.py

@@ -68,6 +68,7 @@ def compile_model(model: GreenFormer, shapeless: bool = False) -> GreenFormer:
     same resolution. Shapeless compile is experimental — the backbone uses
     shape-dependent reshapes that may trigger recompilation.
     """
+    model._compiled = True
     model.__call__ = mx.compile(model.__call__, shapeless=shapeless)  # type: ignore[method-assign]
     return model
 

src/corridorkey_mlx/model/corridorkey.py

@@ -6,6 +6,7 @@ All internal operations use NHWC layout.
 
 from __future__ import annotations
 
+import gc
 from typing import TYPE_CHECKING
 
 import mlx.core as mx
@@ -42,6 +43,7 @@ class GreenFormer(nn.Module):
         self._compute_dtype = dtype
         self._fused_decode = fused_decode
         self._slim = slim
+        self._compiled = False
         self.backbone = HieraBackbone(img_size=img_size)
         self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1)
         self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3)
@@ -69,6 +71,13 @@ class GreenFormer(nn.Module):
         # Backbone always runs in fp32
         features = self.backbone(x)
 
+        # Materialize backbone output so MLX can free intermediate graph nodes.
+        # NOTE: mx.eval is MLX array materialization, not Python eval()
+        if not self._compiled:
+            mx.eval(features)  # noqa: S307
+            gc.collect()
+            mx.clear_cache()
+
         # Cast features to compute dtype for decoders (bf16 saves memory)
         if self._compute_dtype != mx.float32:
             features = [f.astype(self._compute_dtype) for f in features]
@@ -93,6 +102,13 @@ class GreenFormer(nn.Module):
         alpha_coarse = mx.sigmoid(alpha_logits_up)
         fg_coarse = mx.sigmoid(fg_logits_up)
 
+        # Materialize decoder output so MLX can free decoder graph nodes.
+        # NOTE: mx.eval is MLX array materialization, not Python eval()
+        if not self._compiled:
+            mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up)  # noqa: S307
+            gc.collect()
+            mx.clear_cache()
+
         # Refiner receives fp32 coarse predictions
         rgb = x[:, :, :, :3]  # (B, H, W, 3)
         coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1)  # (B, H, W, 4)

tests/test_compilation.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import mlx.core as mx
 import pytest
 
+from corridorkey_mlx.inference.pipeline import compile_model
 from corridorkey_mlx.model.corridorkey import GreenFormer
 
 IMG_SIZE = 256
@@ -32,11 +33,11 @@ def dummy_input() -> mx.array:
 def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> None:
     """Fixed-shape compiled output matches eager output within tolerance."""
     eager_out = model(dummy_input)
-    mx.eval(eager_out)  # noqa: S307
+    mx.eval(eager_out)  # noqa: S307  # NOTE: MLX materialization, not Python eval
 
-    compiled_fn = mx.compile(model.__call__)
-    compiled_out = compiled_fn(dummy_input)
-    mx.eval(compiled_out)  # noqa: S307
+    compiled_model = compile_model(model)
+    compiled_out = compiled_model(dummy_input)
+    mx.eval(compiled_out)  # noqa: S307  # NOTE: MLX materialization
 
     for key in OUTPUT_KEYS:
         diff = float(mx.max(mx.abs(eager_out[key] - compiled_out[key])))
@@ -45,12 +46,12 @@ def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> No
 
 def test_compiled_deterministic(model: GreenFormer, dummy_input: mx.array) -> None:
     """Compiled model produces identical results across consecutive calls."""
-    compiled_fn = mx.compile(model.__call__)
+    compiled_model = compile_model(model)
 
-    out1 = compiled_fn(dummy_input)
-    mx.eval(out1)  # noqa: S307
-    out2 = compiled_fn(dummy_input)
-    mx.eval(out2)  # noqa: S307
+    out1 = compiled_model(dummy_input)
+    mx.eval(out1)  # noqa: S307  # NOTE: MLX materialization
+    out2 = compiled_model(dummy_input)
+    mx.eval(out2)  # noqa: S307  # NOTE: MLX materialization
 
     for key in OUTPUT_KEYS:
         diff = float(mx.max(mx.abs(out1[key] - out2[key])))