Fix DeepSeek-V4 Metal residency leak: materialize per-layer cache state each forward

[snagnever]

Fixes the Metal residency leak in the DeepSeek‑V4 port that aborts generation with [metal::malloc] Resource limit (499000) exceeded after ~11K decode tokens on Apple Silicon (full root cause: ml-explore#1332).

Targeting: DeepSeek‑V4 isn't in main yet, so this PR targets the ml-explore#1192 head branch pc/add-deepseekv4flash-model (which is where mlx_lm/models/deepseek_v4.py lives). @Blaizzy — happy to adjust placement if you'd prefer to fold it into ml-explore#1192 directly.

Cause (short)

The attention caches build per‑decode‑step lazy graphs that are never detached: PoolingCache grows via mx.concatenate, RotatingKVCache via sliced assignment (and the Batch* variants likewise). Because the cache object holds the head of the chain, MLX keeps every prior step's intermediate array — and its backing Metal buffer — resident. The count of live buffers grows ~one per layer per step until it hits the device residency cap (resource_limit, 499000), ≈ 11.3K tokens regardless of prompt length, after which the command queue wedges until a process restart.

Change

Materialize all per‑layer cache state once per forward pass in DeepseekV4Model.__call__, which detaches the chains so the resident‑buffer count stays bounded. One mx.eval per forward (not per cache), so the cost is negligible.

        for layer, layer_cache in zip(self.pipeline_layers, cache):
            h = layer(h, mask, layer_cache, inputs)

        # DeepSeek-V4 Metal OOM fix: materialize all per-layer cache state each
        # forward pass. PoolingCache (concatenate-grow) and RotatingKVCache
        # (sliced assignment), single and batched, otherwise keep each prior
        # step's intermediate arrays -- and their Metal buffers -- resident, so
        # num_resources_ climbs to resource_limit (499000) after ~11.3K tokens.
        # One eval per forward keeps the live-buffer count bounded.
        if cache is not None:
            _cache_arrays = []
            for _c in cache:
                for _leaf in (getattr(_c, "caches", None) or (_c,)):
                    if _leaf is None:
                        continue
                    for _v in vars(_leaf).values():
                        if isinstance(_v, mx.array):
                            _cache_arrays.append(_v)
            if _cache_arrays:
                mx.eval(*_cache_arrays)

Verification (M4 Max 128 GB, mlx 0.31.2)

	before	after
in‑process leak slope	~205 KB/step (linear, no plateau)	7 KB/step
forced 20K‑token generation	OOM at 11,314	clean to 19,989, 0 OOMs
throughput (long gen)	~31 tok/s	31.3 tok/s (no regression)
40‑request tool‑call sweep, single long‑lived server	49 OOMs, aborted at request 20	0 OOMs, 40/40, 19.8 min
knowledge‑bench soak (MMLU+GPQA+HumanEval, 300 requests), single long‑lived server	—	0 OOMs, 300/300, 0 errors, ~2h44m

Notes

• Walking each leaf cache's array attributes (rather than a .state accessor) covers PoolingCache/RotatingKVCache/BatchPoolingCache/BatchRotatingKVCache uniformly; a cache‑class‑local fix is possible but is whack‑a‑mole across four classes (two shared with other models).
• Open question for maintainers: is it intended that realizing the output doesn't detach intermediate cache arrays? A cleaner long‑term fix might be storing the pooled cache with step‑allocated in‑place writes (like KVCache/RotatingKVCache._update_in_place) instead of concatenate, removing the chain at the source.
• (Unrelated heads‑up, not part of this fix: the 2‑bit DQ checkpoint degenerates/loops in open‑ended chat — a quantization‑quality issue independent of this leak. With this fix, such a loop merely runs to max_tokens instead of crashing.)