Models that discard `update_and_fetch`'s values return leak one Metal buffer per layer per generated token — deterministic crash at Metal's buffer-count limit (deepseek_v4 dies at ~11.5k generated tokens)

[aidiffuser]

Summary

GenerationBatch never evaluates cache state during decode (only during
chunked prefill). Any cache field that is not consumed by the logits graph
therefore accumulates an unevaluated lazy slice_update chain — one node
per step — and each node pins that step's freshly materialized input buffer.
The bytes are tiny; the Metal buffer-object count is not. On a 43-layer
model leaking one buffer per layer per token, the process hits Metal's
resource_limit (499,000 buffer objects on M3 Ultra) at ~11,470 generated
tokens and dies with:

RuntimeError: [metal::malloc] Resource limit (499000) exceeded.

This message is easy to misread as a byte-size failure (the same hardware's
max recommended working set is ~499,000 MB — an unfortunate coincidence); it
is a buffer-count failure. All byte metrics (get_active_memory,
get_cache_memory, get_peak_memory) look healthy right up to the crash.

deepseek_v4 is a shipping first-party instance: its attention uses the
K==V trick (SDPA is fed the same rows for keys and values), so both call
sites discard the values half:

win_keys, _ = win_cache.update_and_fetch(k4, k4)   # values return unused

Nothing in the graph ever reads cache.values; at S == 1 the rotating
cache's in-place update rebinds it to slice_update(old_values, new_row)
each step, growing the dead chain.

Observed behavior (production, 2× M3 Ultra 512GB, TP-2)

• Deterministic crash at generated ≈ 11,456–11,488, bracketed to ±8 tokens
  with a per-step probe; independent of prompt length (reproduced with 1k,
  10k, and 23k-token prompts — totals 12.5k/21.7k/34.3k, same generated
  count at death).
• 43 layers × 1 buffer/token: 43 × 11,470 + ~5.8k baseline ≈ 499,000.
• Prompt tokens do not leak because chunked prefill evaluates
  [c.state for c in prompt_cache] every chunk; only decode steps do.
• Total leaked bytes at death: ~0.5 GB (1 KB rows) — invisible.
• Decode throughput stays flat (the dead chain is never executed).
• Severity beyond the process: dying with ~499k live Metal buffers has
  twice coincided, on our machines, with a GPU-driver wedge severe enough
  that WindowServer failed its checkins and watchdogd panicked the host
  (userspace watchdog timeout: no successful checkins from WindowServer).

Minimal reproduction (no model needed)

import mlx.core as mx
from mlx_lm.models.cache import RotatingKVCache

cache = RotatingKVCache(max_size=128)
row = mx.zeros((1, 1, 1, 512), mx.bfloat16)
step = 0
while True:
    keys, _values = cache.update_and_fetch(row, row)  # values discarded
    mx.eval(keys)  # logits-graph analog: keys consumed, values never
    step += 1
    if step % 50_000 == 0:
        print(step, mx.get_active_memory() // 1024, "KB active")
# dies with "[metal::malloc] Resource limit (...) exceeded." when the
# unevaluated cache.values chain pins enough buffers to hit the cap
# (~499k steps here; divide by n_layers for a real model's token count).

Adding cache.values = cache.keys (or mx.eval(cache.values)) after each
update makes it run forever at flat memory.

Fixes

Model-side (what we shipped for deepseek_v4, live-validated through a
13,000-token generation, byte-identical outputs since K == V):

win_keys, _ = win_cache.update_and_fetch(k4, k4)
win_cache.values = win_cache.keys   # drop the dead chain

Library-side options, in rough preference order:

1. GenerationBatch._step (and the non-batch loop) periodically evaluate
   cache state during decode — e.g. fold c.state into the existing
   async_eval every N steps. This closes the whole bug class for every
   model, not just K==V ones, at negligible cost.
2. A K-only cache variant (or an update_and_fetch(keys) overload) for
   models that never store distinct values.
3. At minimum, a documentation note: cache fields not consumed by the
   returned graph must be evaluated or rebound by the model, or they
   accumulate unevaluated updates without bound.

Related context: this is the second silent-degradation mode we have hit in
this area (ml-explore/mlx#3964 covers CompilerCache growth from int
constants; #1189 the deepseek_v4 loader fixes). We are running the
call-site fix in production and can validate a library-side patch on the
same workload.

Co-authored with Claude Fable 5 (Anthropic)