corridorkey-mrp-mlx/docs/plans/2026-03-01-feat-2048-smoke-test-plan.md

title	type	date
feat: Add 2048 smoke test	feat	2026-03-01

Add 2048 Smoke Test

Overview

Validate MLX model runs end-to-end at CorridorKey's native 2048×2048 resolution on Apple Silicon. Smoke/stability check — not a new parity campaign.

Problem Statement

All existing tests/scripts default to 256–512. No automated path confirms 2048 actually works. The model was trained at 2048, so pos_embed interpolation is a no-op at that resolution, but we haven't exercised the full pipeline there.

Proposed Solution

One new script + one skipable pytest test. Reuse CorridorKeyMLXEngine (already defaults to img_size=2048). Minimal additions.

Deliverables

1. scripts/smoke_2048.py

Mirrors scripts/smoke_engine.py pattern but targeted at 2048 with richer diagnostics.

CLI args:

• --checkpoint PATH (default: checkpoints/corridorkey_mlx.safetensors)
• --img-size INT (default: 2048 — this script's whole point)
• --image PATH (optional — uses synthetic input if omitted)
• --hint PATH (optional — uses synthetic hint if omitted)
• --output-dir DIR (default: output/smoke_2048/)
• --save-outputs / --no-save-outputs (default: save)
• --seed INT (default: 42)
• --compile / --no-compile (default: True — engine default)

Behavior:

1. Generate or load 2048×2048 RGB + hint inputs
2. Instantiate CorridorKeyMLXEngine(checkpoint_path, img_size, compile)
3. Reset peak memory via mx.metal.reset_peak_memory() (if available, try/except)
4. Time engine.process_frame(rgb, mask) with wall-clock
5. Read peak memory via mx.metal.get_peak_memory() (if available)
6. Report: input shape, output shapes, dtypes, elapsed time, peak memory (MB)
7. Diagnostic: min/max/mean of alpha+fg, NaN/Inf check, flag all-zeros/all-ones
8. Optionally save alpha.png, fg.png, comp.png
9. Catch RuntimeError/MemoryError with actionable error message

Synthetic input generation:

rng = np.random.default_rng(seed)
rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8)
# Circular gradient hint — exercises the path better than random noise
mask = <simple radial gradient uint8>

2. tests/test_smoke_2048.py

Two tests:

1. test_smoke_2048_full — loads checkpoint, runs engine at 2048, asserts shapes + no NaN. Marked @pytest.mark.skipif(not HAS_CHECKPOINT) AND @pytest.mark.slow. Won't run in normal uv run pytest; run via uv run pytest -m slow.

2. test_smoke_2048_wiring — lightweight, no checkpoint. Verifies GreenFormer(img_size=2048) constructs OK and produces correct output shapes with random weights at 2048 (or smaller like 256 if too heavy without checkpoint). Runs in normal suite.

3. README update

Add "2048 Smoke Test" section after the existing "Smoke test" subsection:

### 2048 smoke test

Validates full end-to-end inference at CorridorKey's native 2048 resolution.
Uses synthetic inputs by default — no test images required.

```bash
uv run python scripts/smoke_2048.py

To use real images:

uv run python scripts/smoke_2048.py --image shot.png --hint hint.png

Reports timing, peak memory, output shapes, and value-range diagnostics. This is an execution check, not a 2048 parity validation.

## Technical Considerations

- **Memory**: 2048×2048×4 float32 ≈ 64MB input, but backbone + decoder intermediates will be much larger. Engine already handles this; we just report peak.
- **Compile**: Engine defaults `compile=True`. First call at 2048 will be slow (compile cost). The smoke script is a single-run tool so the compile overhead is included in timing.
- **Pos_embed**: At 2048, no interpolation needed (trained at 2048). This is the easiest resolution to test.
- **`mx.metal`**: May not be available in all environments. Wrap in try/except.

## Acceptance Criteria

- [x] `scripts/smoke_2048.py` runs successfully with checkpoint
- [x] Reports: input size, output shapes, elapsed time, peak memory
- [x] Detects NaN/Inf and flags suspicious outputs
- [x] Supports both synthetic and user-supplied inputs
- [x] `test_smoke_2048_full` passes when checkpoint present + `-m slow`
- [x] `test_smoke_2048_wiring` passes in normal test suite
- [x] README documents the smoke test
- [x] Existing 512 tests unchanged (112 pass, 1 slow deselected)
- [x] No new binary artifacts committed

## Files Changed

| File | Change |
|------|--------|
| `scripts/smoke_2048.py` | New — main smoke script |
| `tests/test_smoke_2048.py` | New — pytest smoke + wiring tests |
| `README.md` | Add "2048 smoke test" subsection |

## Unresolved Questions

- Exact peak memory at 2048 unknown until first run — might be tight on 8GB machines?
- Include `--compile` timing breakdown (warmup vs inference) or just total?