feat: slim forward mode — drop intermediate tensor refs

slim=True returns 4-key dict, engine uses it, pipeline.infer keeps full.
2 new tests: key count + bit-exact match vs full output.

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

src/corridorkey_mlx/inference/pipeline.py

@@ -35,6 +35,7 @@ def load_model(
     shapeless: bool = False,
     dtype: mx.Dtype = mx.bfloat16,
     fused_decode: bool = True,
+    slim: bool = False,
 ) -> GreenFormer:
     """Build GreenFormer and load weights from safetensors checkpoint.
 
@@ -50,8 +51,10 @@ def load_model(
             backbone and sigmoid always stay fp32. All outputs are fp32.
         fused_decode: If True, batch alpha+fg decoder upsamples to reduce
             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.
     """
-    model = GreenFormer(img_size=img_size, dtype=dtype, fused_decode=fused_decode)
+    model = GreenFormer(img_size=img_size, dtype=dtype, fused_decode=fused_decode, slim=slim)
     model.load_checkpoint(checkpoint)
     if compile:
         model = compile_model(model, shapeless=shapeless)

src/corridorkey_mlx/model/corridorkey.py

@@ -36,10 +36,12 @@ class GreenFormer(nn.Module):
         img_size: int = 512,
         dtype: mx.Dtype = mx.float32,
         fused_decode: bool = False,
+        slim: bool = False,
     ) -> None:
         super().__init__()
         self._compute_dtype = dtype
         self._fused_decode = fused_decode
+        self._slim = slim
         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)
@@ -100,6 +102,14 @@ class GreenFormer(nn.Module):
         alpha_final = mx.sigmoid(alpha_logits_up + delta_logits[:, :, :, 0:1])
         fg_final = mx.sigmoid(fg_logits_up + delta_logits[:, :, :, 1:4])
 
+        if self._slim:
+            return {
+                "alpha_coarse": alpha_coarse,
+                "fg_coarse": fg_coarse,
+                "alpha_final": alpha_final,
+                "fg_final": fg_final,
+            }
+
         return {
             "alpha_logits": alpha_logits.astype(mx.float32),
             "fg_logits": fg_logits.astype(mx.float32),

tests/test_model_contract.py

@@ -290,3 +290,49 @@ def test_fp32_default_unchanged() -> None:
             np.array(out_explicit[key]),
             err_msg=f"{key} differs between default and explicit fp32",
         )
+
+
+# ---------------------------------------------------------------------------
+# Slim forward mode
+# ---------------------------------------------------------------------------
+
+SLIM_KEYS = {"alpha_coarse", "fg_coarse", "alpha_final", "fg_final"}
+
+
+def test_slim_returns_four_keys() -> None:
+    """slim=True returns only the 4 final/coarse keys."""
+    model = GreenFormer(img_size=SMALL_IMG_SIZE, slim=True)
+    x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
+    out = model(x)
+    # NOTE: mx.eval is MLX array materialization, not Python eval()
+    mx.eval(out)  # noqa: S307
+    assert set(out.keys()) == SLIM_KEYS
+
+
+def test_slim_matches_full_output() -> None:
+    """slim=True values match the corresponding keys from full output."""
+    from mlx.utils import tree_flatten
+
+    mx.random.seed(55)
+    x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4))
+    # NOTE: mx.eval is MLX array materialization, not Python eval()
+    mx.eval(x)  # noqa: S307
+
+    model_full = GreenFormer(img_size=SMALL_IMG_SIZE, slim=False)
+    mx.eval(model_full.parameters())  # noqa: S307
+
+    model_slim = GreenFormer(img_size=SMALL_IMG_SIZE, slim=True)
+    model_slim.load_weights(tree_flatten(model_full.parameters()))  # type: ignore[arg-type]
+    mx.eval(model_slim.parameters())  # noqa: S307
+
+    out_full = model_full(x)
+    out_slim = model_slim(x)
+    mx.eval(out_full)  # noqa: S307
+    mx.eval(out_slim)  # noqa: S307
+
+    for key in SLIM_KEYS:
+        np.testing.assert_array_equal(
+            np.array(out_full[key]),
+            np.array(out_slim[key]),
+            err_msg=f"{key} differs between slim and full forward",
+        )