Fix memory leak: close BatchGenerator properly and clear Metal cache

[janhilgard]

Summary

• Fix VRAM (wired memory) growth during generation by properly closing BatchGenerator and clearing Metal cache periodically
• Without this fix, wired memory spikes from ~70GB to 130GB+ during generation on large contexts
• With fix: stable ~108GB during generation, returns to baseline within ~10s after completion
• KV cache is preserved for agentic multi-turn reuse

Root Cause

Three issues caused unbounded VRAM growth:

1. BatchGenerator.close() never called — when the scheduler replaced a BatchGenerator, it set self.batch_generator = None without calling .close(), leaving wired_limit elevated and Metal buffers unreleased

2. No periodic cache clearing — MLX accumulates command buffers and intermediate tensors during generation. Without explicit mx.clear_cache(), wired memory grows with every generated token

3. No engine-level cleanup — after the scheduler finishes processing requests, the engine core never freed Metal buffers

Changes

File	Change
scheduler.py	Add _close_batch_generator() — calls .close() before discarding
scheduler.py	mx.clear_cache() every 32 scheduler steps (key: limits VRAM spike during generation)
scheduler.py	mx.clear_cache() when requests finish in _cleanup_finished
scheduler.py	mx.eval() + mx.clear_cache() after KV cache extraction (prevents deferred spike)
scheduler.py	Use _close_batch_generator() in reset() and _recover_from_cache_error()
scheduler.py	Add Metal memory stats (active/peak/cache GB) to get_stats()
engine_core.py	mx.clear_cache() after distributing finished request outputs
tests/test_memory_stability.py	Unit tests for all memory management paths

Test Results

Qwen3-Next-80B-A3B-6bit on 256GB M3 Ultra, ~5K token contexts:

During generation (wired memory):

Without fix	With fix
70 → 132 GB (spike)	107 → 109 GB (stable)

Multi-request stability (5 sequential requests):

Request	Before	Peak (during gen)	+10s after
1	45.32 GB	106.45 GB	45.67 GB
2	45.67 GB	106.57 GB	45.67 GB
3	45.67 GB	106.68 GB	45.67 GB
4	45.67 GB	106.81 GB	45.67 GB
5	45.67 GB	106.92 GB	30.37 GB

Memory returns to baseline after each request — no monotonic growth.

Test plan

• Unit tests pass (tests/test_memory_stability.py)
• VRAM spike during generation reduced from +62GB to +2GB
• Multi-request: wired memory bounded, returns to baseline
• KV cache properly stored for agentic multi-turn
• Throughput unchanged (~62-64 tok/s)
• CI passes

🤖 Generated with Claude Code

[janhilgard]

Closing — all changes from this PR were already merged into main via PR #46 (b191aec).