`ArraysCache.advance()` still leaks one Metal buffer per SSM layer per token through `BatchGenerator` after #1632 — long completions crash at ~499000/(#SSM−1) tokens on main

[aidiffuser]

Summary

#1632 made ArraysCache.state carry left_padding / lengths, so the lazy chain that advance() builds is bounded wherever .state is evaluated. That covers prompt processing, but GenerationBatch._step never evaluates cache state — it only evaluates the sampled tokens and logprobs (generate.py L1430-L1435). The hybrid models rebuild the SSM mask from cache[ssm_idx] only (qwen3_next L415, qwen3_5 L289, kimi_linear L451), so on every other GatedDeltaNet/KDA layer the left_padding -= N in advance() stays an unevaluated x - 1 - 1 - … chain for the whole completion. Each node pins that step's scalar constant as a live Metal buffer, and Metal counts objects, not bytes (mx.device_info()["resource_limit"] = 499000 here), so any single completion longer than about

499000 / (num_ssm_layers - 1)   tokens

kills the process with [metal::malloc] Resource limit (499000) exceeded — while get_active_memory() stays flat (35 layers × 4 B ≈ 140 B/token), which is what makes it look like anything but a leak.

This is the same site as #1641 (closed by #1632). The regression test added in #1632 calls mx.eval(cache.state) after every advance(), which is exactly the step the batched decode loop does not perform, so the test passes while mlx_lm.server still crashes. #1780 / #1784 proposed evaluating cache state periodically and were declined as a workaround, so below is a root-cause fix that removes the graph instead.

Repro on main @ 14e89f1 (mlx 0.32.2, M3 Ultra) — no downloads, ~3 min to the crash

Tiny random-weight qwen3_next (48 layers → 36 linear-attention layers, a few MB), one sequence decoded through BatchGenerator, i.e. the server path:

$ python repro_arrays_cache_advance_batch.py
mlx 0.32.2, resource_limit=499000, layers=48 linear=36, predicted crash ~14257 tokens, patched=False
  2000 tokens, active=20.2 MB, 23s
  4000 tokens, active=26.5 MB, 47s
  6000 tokens, active=32.8 MB, 71s
  8000 tokens, active=39.0 MB, 95s
  10000 tokens, active=45.3 MB, 119s
  12000 tokens, active=50.9 MB, 144s
  14000 tokens, active=57.1 MB, 170s
CRASH after 14215 generated tokens (173s): [metal::malloc] Resource limit (499000) exceeded.

$ python repro_arrays_cache_advance_batch.py --patched
mlx 0.32.2, resource_limit=499000, layers=48 linear=36, predicted crash ~14257 tokens, patched=True
  2000 tokens, active=19.9 MB, 23s
  4000 tokens, active=25.9 MB, 46s
  6000 tokens, active=32.0 MB, 68s
  8000 tokens, active=38.0 MB, 90s
  10000 tokens, active=44.0 MB, 114s
  12000 tokens, active=49.3 MB, 137s
  14000 tokens, active=55.3 MB, 160s
  16000 tokens, active=61.3 MB, 182s
  18000 tokens, active=67.3 MB, 204s
  20000 tokens, active=73.3 MB, 227s
OK: 20000 tokens generated, active=73.3 MB (227s)

repro_arrays_cache_advance_batch.py

"""Repro: ArraysCache.advance() still leaks one Metal buffer per SSM layer per decoded token
through BatchGenerator on mlx-lm main (post-#1632), for any completion long enough.

Builds a tiny random-weight qwen3_next (48 layers -> 36 GatedDeltaNet layers, a few MB) and
decodes ONE sequence through BatchGenerator, i.e. the mlx_lm.server path. No downloads.

    python repro_arrays_cache_advance_batch.py            # dies: [metal::malloc] Resource limit (499000) exceeded
    python repro_arrays_cache_advance_batch.py --patched  # same run with advance() deferred to a Python int: flat

Predicted crash: 499000 / (36 - 1) ~= 14,257 generated tokens (only cache[ssm_idx]'s mask is
ever rebuilt, so the other 35 layers' `left_padding - 1 - 1 ...` chains are never evaluated).
"""
import sys, time
import mlx.core as mx
from mlx_lm.generate import BatchGenerator
from mlx_lm.models import qwen3_next
from mlx_lm.models.cache import ArraysCache

PATCHED = "--patched" in sys.argv
MAX_TOKENS = 20000

if PATCHED:
    # Proposed root-cause fix: advance() accumulates a Python int; left_padding / lengths fold it
    # in on read. No per-step graph, identical values, nothing to evaluate.
    def _get(raw, pending):
        def getter(self):
            value = getattr(self, raw, None)
            n = getattr(self, pending, 0)
            if value is not None and n:
                value = value - n
                setattr(self, raw, value)
                setattr(self, pending, 0)
            return value
        return getter

    def _set(raw, pending):
        def setter(self, value):
            setattr(self, raw, value)
            setattr(self, pending, 0)
        return setter

    ArraysCache.left_padding = property(_get("_raw_left_padding", "_pending_left_padding"), _set("_raw_left_padding", "_pending_left_padding"))
    ArraysCache.lengths = property(_get("_raw_lengths", "_pending_lengths"), _set("_raw_lengths", "_pending_lengths"))

    def advance(self, N):
        if getattr(self, "_raw_lengths", None) is not None:
            self._pending_lengths = getattr(self, "_pending_lengths", 0) + N
        if getattr(self, "_raw_left_padding", None) is not None:
            self._pending_left_padding = getattr(self, "_pending_left_padding", 0) + N

    ArraysCache.advance = advance

args = qwen3_next.ModelArgs(
    model_type="qwen3_next", hidden_size=64, num_hidden_layers=48, intermediate_size=128,
    num_attention_heads=2, linear_num_value_heads=2, linear_num_key_heads=1,
    linear_key_head_dim=16, linear_value_head_dim=16, linear_conv_kernel_dim=4,
    num_experts=4, num_experts_per_tok=2, decoder_sparse_step=1,
    shared_expert_intermediate_size=64, mlp_only_layers=[], moe_intermediate_size=64,
    rms_norm_eps=1e-6, vocab_size=256, num_key_value_heads=1, rope_theta=10000.0,
    partial_rotary_factor=0.25, max_position_embeddings=65536, head_dim=32,
)
model = qwen3_next.Model(args)
mx.eval(model.parameters())
n_linear = sum(l.is_linear for l in model.model.layers)
print(f"mlx {mx.__version__}, resource_limit={mx.device_info().get('resource_limit')}, "
      f"layers={args.num_hidden_layers} linear={n_linear}, predicted crash ~{499000 // (n_linear - 1)} tokens, patched={PATCHED}", flush=True)

gen = BatchGenerator(model, max_tokens=MAX_TOKENS, prefill_step_size=512)
gen.insert([[1, 2, 3, 4, 5, 6, 7, 8]], max_tokens=[MAX_TOKENS])
generated, t0 = 0, time.perf_counter()
try:
    while True:
        responses = gen.next_generated()   # BatchGenerator.next() returns (prompt_responses, generation_responses)
        if not responses:
            break
        generated += len(responses)
        if generated % 2000 == 0:
            print(f"  {generated} tokens, active={mx.get_active_memory() / 2**20:.1f} MB, "
                  f"{time.perf_counter() - t0:.0f}s", flush=True)
        if any(getattr(r, "finish_reason", None) for r in responses):
            break
except RuntimeError as error:
    print(f"CRASH after {generated} generated tokens ({time.perf_counter() - t0:.0f}s): {error}")
    sys.exit(1)
print(f"OK: {generated} tokens generated, active={mx.get_active_memory() / 2**20:.1f} MB ({time.perf_counter() - t0:.0f}s)")

The active= column grows identically in both runs — that is the 12 attention layers' KV cache. Byte telemetry cannot see this leak; only the buffer-object count does.

Model-free version of the same thing (the #1632 test loop without the per-step mx.eval(cache.state); 34 merged caches, only cache[0]'s mask consumed):

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

caches = [ArraysCache.merge([ArraysCache(2)]) for _ in range(34)]
for step in range(1, 20001):
    mx.eval(caches[0].make_mask(1))       # the model consumes one layer's mask
    for c in caches:
        c.advance(1)                      # every SSM layer advances
# -> RuntimeError: [metal::malloc] Resource limit (499000) exceeded. at step 15121 (= 499000 // 33), ~9 s

Production data point

GLM-5.3-Flash (34 KDA layers, same ArraysCache pattern) served batched on a 2× M3 Ultra cluster died at exactly 15,121 generated tokens of one completion — 499000 / 33 — after having served ~90k tokens across shorter turns without incident. The cap is per completion, not per process, so it only shows up on long reasoning streams.

Proposed root-cause fix

advance() never needs a graph: accumulate the shift as a Python int and fold it into left_padding / lengths when they are read (make_mask, filter, extend, state). Identical values, no per-step arrays, no synchronization:

class ArraysCache(_BaseCache):
    def __new__(cls, *args, **kwargs):
        instance = super().__new__(cls)
        instance._left_padding = None
        instance._lengths = None
        instance._pending = 0          # tokens advanced since the arrays were last folded
        return instance

    def _fold(self):
        if self._pending:
            if self._lengths is not None:
                self._lengths = self._lengths - self._pending
            if self._left_padding is not None:
                self._left_padding = self._left_padding - self._pending
            self._pending = 0

    @property
    def left_padding(self):
        self._fold()
        return self._left_padding

    @left_padding.setter
    def left_padding(self, value):
        self._fold()
        self._left_padding = value

    # same property pair for `lengths`

    def advance(self, N):
        if self._lengths is not None or self._left_padding is not None:
            self._pending += N

Every existing call site keeps reading cache.left_padding / cache.lengths unchanged, and the one layer whose mask is consumed still evaluates a single subtraction per step. We run this (as a runtime patch on top of 0.31.3) in production; the --patched run above is this change monkeypatched onto main. Happy to turn it into a PR with a regression test that asserts no -> edges after N un-evaluated advance() calls, if you want it.

Co-authored with Claude Fable 5.1