#!/usr/bin/env python3
"""Smoking-gun repro: qwen3_5/qwen3_next decode leaks live Metal buffer OBJECTS.
pip install mlx mlx-lm
python repro_leak.py # hybrid: dies with
# RuntimeError: [metal::malloc]
# Resource limit (499000) exceeded
# after ~10.4k decode tokens, minutes
python repro_leak.py --attention-only # same depth, no SSM: flat, no crash
One script, one knob. The model is a random-weight qwen3_5_moe built in
seconds (hidden 64, vocab 128, ~MBs of weights — NO download): the leak rate
scales with the number of GatedDeltaNet (SSM) layers, not model size, so a
DEEP-but-tiny model (default: 50 layers = 49 SSM + 1 attention) hits the real
crash quickly on a stock pip install.
What leaks
----------
`ArraysCache.advance()` decrements `left_padding`/`lengths` with lazy
mx.array arithmetic once per SSM layer per decoded token
(mlx_lm/models/qwen3_next.py:302 -> mlx_lm/models/cache.py:685). The model
only rebuilds the SSM mask from cache[ssm_idx=0] (qwen3_next.py:401,415), so
for every OTHER SSM layer the graph chain is dead: never evaluated, one new
node per token, each node pinning a fresh scalar array whose live 4-byte
MTL::Buffer counts against the per-process resource limit
(mlx/backend/metal/allocator.cpp:142, 499000 on Apple silicon). Bytes stay
flat; the object count climbs at ~(#SSM layers - 1)/token; `mx.clear_cache()`
cannot help (live handles, not cache). The server/batch path arms this even
at batch size 1 (`ArraysCache.merge` sets left_padding, generate.py:1036 /
cache.py:709); attention layers have equivalent per-step array state
(`BatchKVCache.offset`) but RoPE consumes it every step, which is why vanilla
transformers (gpt-oss etc.) do not leak.
Observables
-----------
- stock mlx: deterministic crash at ~499000/(#SSM-1) decode tokens; monotonic
process-RSS growth (chain nodes) and a slow monotonic
`mx.get_active_memory()` creep (~4 B/object) while cache memory is flat.
- probe build (mlx + buffer-count probes: get_active_buffer_count et al.):
exact live-object slope per token, and a per-size-class histogram pinning
the growth to the smallest size class.
- with the ArraysCache fix (mlx-lm patch): slope 0, no crash, same tokens.
Exit codes: 0 flat/survived, 1 leak detected (probe build), 2 resource-limit
crash reproduced.
"""
from __future__ import annotations
import argparse
import json
import resource
import sys
import time
import mlx.core as mx
def _mx_fn(*names):
"""First available mx function: upstream buffer-count probe naming first,
then the name used by the local v0.32.0 probe wheel."""
for n in names:
if hasattr(mx, n):
return getattr(mx, n)
return None
_active_count = _mx_fn("get_active_buffer_count", "get_active_resource_count")
_cache_count = _mx_fn("get_cache_buffer_count", "get_cache_count")
_histogram = _mx_fn("get_buffer_histogram", "get_resource_histogram")
HAVE_COUNT_PROBE = _active_count is not None
def build_model(num_layers: int, interval: int, experts: int, hidden: int = 64):
from mlx_lm.models import qwen3_5_moe
cfg = dict(
model_type="qwen3_5_moe",
hidden_size=hidden,
intermediate_size=hidden * 2,
num_hidden_layers=num_layers,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=16,
vocab_size=128,
full_attention_interval=interval,
linear_num_value_heads=4,
linear_num_key_heads=2,
linear_key_head_dim=16,
linear_value_head_dim=16,
linear_conv_kernel_dim=4,
num_experts=experts,
num_experts_per_tok=2,
moe_intermediate_size=32,
shared_expert_intermediate_size=32,
)
args = qwen3_5_moe.ModelArgs.from_dict(
{"model_type": "qwen3_5_moe", "text_config": cfg}
)
model = qwen3_5_moe.Model(args)
mx.eval(model.parameters())
n_ssm = sum(1 for l in model.layers if getattr(l, "is_linear", False))
print(
f"[model] layers={num_layers} (ssm={n_ssm} attn={num_layers - n_ssm}) "
f"interval={interval} experts={experts} hidden={hidden}",
flush=True,
)
return model, n_ssm
def rss_mb() -> float:
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 # macOS: bytes
def probe(sync: bool = True) -> dict:
"""Sample allocator state. Synchronizes first so in-flight command buffers
(whose completion handlers release temporaries) have drained."""
if sync:
mx.synchronize()
out = {
"active_mb": mx.get_active_memory() / 1e6,
"cache_mb": mx.get_cache_memory() / 1e6,
"rss_mb": rss_mb(),
}
if HAVE_COUNT_PROBE:
out["count"] = _active_count()
out["cache_count"] = _cache_count()
# Live handles are what crash the process; cached objects count
# against the limit too but are reclaimable (the allocator clears the
# cache under count pressure, and byte limits bound its growth).
out["live"] = out["count"] - out["cache_count"]
out["hist"] = dict(_histogram())
return out
def fmt(p: dict) -> str:
base = (
f"active={p['active_mb']:9.2f}MB cache={p['cache_mb']:7.1f}MB "
f"rss={p['rss_mb']:7.0f}MB"
)
if HAVE_COUNT_PROBE:
return (
f"live={p['live']:>7d} (+cached={p['cache_count']} "
f"= {p['count']}) " + base
)
return base + " (stock mlx: no count probe)"
def hist_delta(a: dict, b: dict, top: int = 8) -> str:
keys = sorted(set(a.get("hist", {})) | set(b.get("hist", {})))
deltas = [
(k, b.get("hist", {}).get(k, 0) - a.get("hist", {}).get(k, 0)) for k in keys
]
deltas = [d for d in deltas if d[1] != 0]
deltas.sort(key=lambda kv: -abs(kv[1]))
if not deltas:
return " (no size-class drift)"
return "\n".join(f" size<= {k:>10d} B : {d:+d}" for k, d in deltas[:top])
def _cache_arrays(c):
st = getattr(c, "state", None)
if st is None:
st = []
arrs = [a for a in (st if isinstance(st, (list, tuple)) else [st]) if a is not None]
# ArraysCache.state exposes only the state slots; the leaking graph chains
# hang off left_padding/lengths, so include those explicitly.
for attr in ("left_padding", "lengths", "offset"):
v = getattr(c, attr, None)
if isinstance(v, mx.array):
arrs.append(v)
return arrs
def run_plain(model, n_tokens, report_every, samples, eval_cache_every):
from mlx_lm.models.cache import make_prompt_cache
cache = make_prompt_cache(model)
y = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]])
logits = model(y, cache=cache)
y = mx.argmax(logits[:, -1], axis=-1)
mx.eval(y)
samples.append((0, probe()))
print(f"tok {0:>6d} {fmt(samples[-1][1])}", flush=True)
for t in range(1, n_tokens + 1):
logits = model(y[:, None], cache=cache)
y = mx.argmax(logits[:, -1], axis=-1)
mx.async_eval(y)
if eval_cache_every and t % eval_cache_every == 0:
mx.eval([a for c in cache for a in _cache_arrays(c)])
if t % report_every == 0:
samples.append((t, probe()))
print(f"tok {t:>6d} {fmt(samples[-1][1])}", flush=True)
return samples
def run_batch(model, n_tokens, report_every, samples, eval_cache_every):
"""Drive mlx_lm's BatchGenerator directly -- the exact mlx_lm.server decode
path (left-padded batch caches, masked SSM kernel, async_eval streaming)."""
from mlx_lm.generate import BatchGenerator
bg = BatchGenerator(model, max_tokens=n_tokens + 8)
bg.insert([[1, 2, 3, 4, 5, 6, 7, 8]], max_tokens=[n_tokens])
t = 0
samples.append((0, probe()))
print(f"tok {0:>6d} {fmt(samples[-1][1])}", flush=True)
next_report = report_every
done = False
idle = 0
while not done:
prompt_responses, gen_responses = bg.next()
idle = idle + 1 if not prompt_responses and not gen_responses else 0
if idle > 32:
break
for r in gen_responses:
t += 1
if getattr(r, "finish_reason", None) is not None:
done = True
if eval_cache_every and t >= 1 and t % eval_cache_every == 0:
for entry in bg._generation_batch.prompt_cache:
arrs = _cache_arrays(entry)
if arrs:
mx.eval(arrs)
if t >= next_report:
samples.append((t, probe()))
print(f"tok {t:>6d} {fmt(samples[-1][1])}", flush=True)
next_report += report_every
bg.close()
return samples
def main():
ap = argparse.ArgumentParser(
description="qwen3_5 hybrid-SSM Metal buffer-object leak repro"
)
ap.add_argument("--mode", choices=["plain", "batch"], default="batch")
ap.add_argument(
"--attention-only",
action="store_true",
help="THE control knob: same depth, zero SSM layers -- no leak, no crash",
)
ap.add_argument("--layers", type=int, default=50)
ap.add_argument(
"--interval",
type=int,
default=0,
help="full_attention_interval (0 = auto: == layers for the hybrid "
"smoking gun, 1 with --attention-only). Production qwen3_5 uses 4.",
)
ap.add_argument("--experts", type=int, default=8)
ap.add_argument("--hidden", type=int, default=64)
ap.add_argument(
"--tokens",
type=int,
default=12000,
help="decode budget; the default hybrid config is expected to crash at "
"~499000/(#SSM-1) ~= 10.4k tokens before reaching it",
)
ap.add_argument("--report-every", type=int, default=250)
ap.add_argument("--disable-compile", action="store_true")
ap.add_argument(
"--ops",
action="store_true",
help="force the pure-ops gated delta path instead of the custom metal kernel",
)
ap.add_argument(
"--eval-cache-every",
type=int,
default=0,
help="mx.eval all cache state arrays every N tokens (0 = never). "
"Flattening the count here proves the leak is unevaluated cache-state "
"graph chains.",
)
ap.add_argument("--fail-slope", type=float, default=0.5)
ap.add_argument("--json-out", type=str, default="")
args = ap.parse_args()
interval = args.interval
if interval == 0:
interval = 1 if args.attention_only else args.layers
elif args.attention_only:
interval = 1
if args.disable_compile:
mx.disable_compile()
print("[env] mx.disable_compile()")
if args.ops:
from mlx_lm.models import gated_delta
orig = gated_delta.gated_delta_update
def ops_update(*a, **kw):
kw["use_kernel"] = False
return orig(*a, **kw)
import mlx_lm.models.qwen3_5 as _q35
import mlx_lm.models.qwen3_next as _q3n
for m in (gated_delta, _q3n, _q35):
if hasattr(m, "gated_delta_update"):
m.gated_delta_update = ops_update
print("[env] gated_delta_update forced to ops path (no custom metal kernel)")
mx.random.seed(0)
model, n_ssm = build_model(args.layers, interval, args.experts, args.hidden)
expected = None
if n_ssm >= 2 and args.mode == "batch":
expected = 499_000 // (n_ssm - 1)
print(
f"[predict] leak ~{n_ssm - 1} objects/token -> resource-limit crash "
f"near token ~{expected:,} (if unpatched)",
flush=True,
)
samples: list = []
t0 = time.time()
crashed = None
runner = run_batch if args.mode == "batch" else run_plain
try:
runner(model, args.tokens, args.report_every, samples, args.eval_cache_every)
except RuntimeError as e:
if "Resource limit" not in str(e):
raise
crashed = str(e)
dt = time.time() - t0
last_t = samples[-1][0] if samples else 0
verdict = None
slope = None
leaking = False
if crashed:
leaking = True
verdict = (
f"CRASH REPRODUCED: {crashed!r} after ~{last_t}+ decode tokens "
f"(predicted ~{expected:,}; {dt:.0f}s wall)"
)
elif HAVE_COUNT_PROBE and len(samples) >= 4:
# Slope over the mid-run steady state: skip warmup AND the final
# sample (after the request completes, the batch drops its caches and
# the chains release into the buffer cache -- by design).
half = samples[len(samples) // 4 : -1]
if len(half) < 2:
half = samples[-2:]
(t_a, p_a), (t_b, p_b) = half[0], half[-1]
if t_b > t_a:
slope = (p_b["live"] - p_a["live"]) / (t_b - t_a)
leaking = slope > args.fail_slope
verdict = (
f"LEAKING: +{slope:.2f} LIVE buffer objects/token "
f"(tokens {t_a}..{t_b}: live {p_a['live']} -> {p_b['live']}; "
f"predicted ~{max(n_ssm - 1, 0)})"
if leaking
else f"FLAT: {slope:+.3f} live objects/token (tokens {t_a}..{t_b})"
)
print("\n[size-class drift over second half]")
print(hist_delta(p_a, p_b))
else:
verdict = (
"SURVIVED (stock build, no count probe -- see rss/active byte proxies)"
)
print(f"\n[verdict] {verdict}")
print(f"[run] {last_t} tokens in {dt:.1f}s ({last_t/max(dt,1e-9):.0f} tok/s)")
if args.json_out:
payload = {
"argv": sys.argv[1:],
"n_ssm": n_ssm,
"interval": interval,
"have_count_probe": HAVE_COUNT_PROBE,
"crashed": crashed,
"slope_per_token": slope,
"leaking": leaking,
"expected_crash_token": expected,
"samples": [
(t, {k: v for k, v in p.items() if k != "hist"}) for t, p in samples
],
"hist_first": samples[0][1].get("hist") if samples else None,
"hist_last": samples[-1][1].get("hist") if samples else None,
"seconds": dt,
}
with open(args.json_out, "w") as f:
json.dump(payload, f, indent=2)
print(f"[json] {args.json_out}")
sys.exit(2 if crashed else (1 if leaking else 0))
if __name__ == "__main__":
main()
84-second repro (stock pip install, no downloads)
repro_leak.py(inlined at the bottom) builds a random-weight 50-layer tiny qwen3_5_moe (hidden size 64, a few MB of weights, built in seconds) and decodes throughBatchGenerator— the exactmlx_lm.serverdecode path. The leak rate scales with the number of GatedDeltaNet (SSM) layers, not model size, so a deep-but-tiny model hits the real production crash in minutes. No instrumentation needed to see it on stock mlx: process RSS climbs monotonically 113 MB → 1.44 GB (~130 KB/token of graph-chain descriptors) whilemx.get_active_memory()stays ~20 MB.The measurement
Decoding any qwen3_5 / qwen3_next (hybrid GatedDeltaNet) model through the batch pipeline accumulates live Metal buffer objects — not bytes — at exactly (#SSM layers − 1) per decoded token. Measured on
mlx-community/Qwen3.5-122B-A10B-6bit(36 SSM layers), M5 Max 128 GB, macOS 26.5.2, mlx 0.32.0 + mlx-lm 0.31.3, stock code:Byte memory stays flat (the leaked objects are 4-byte scalars ≈ 140 B/token), so
mx.get_active_memory()/mx.get_cache_memory()see nothing, andmx.clear_cache()cannot help — these are live handles, not cache. The process dies at the per-process object cap (499,000 on Apple silicon) as soon as a single completion exceeds ~499000/(#SSM−1) tokens — ~14.1k for the 122B. For isolated or fully drained batches, requests below that threshold are unaffected (the chains die with the request), which is why this presents as flaky long-generation crashes. (A second, much slower variant of the same pattern accumulates across batch churn — lazy gather/concat on the same unread metadata infilter/extend— so a long-lived continuously batched server is not strictly safe below the threshold either; the fix covers both.)Root cause
ArraysCache.advance()(cache.py#L685 @ v0.31.3; identical on main):left_padding/lengthsare mx.arrays in the batch/server path, so-=is lazy graph arithmetic: each call appends an unevaluatedSubtractnode whose inputs are (previous node, fresh scalar array holding a live 4-byteMTL::Bufferfrom construction). Every GatedDeltaNet layer callscache.advance(S)per forward (qwen3_next.py:302), but the model rebuilds the SSM mask from only one of those caches —ssm_idx = 0(qwen3_next.py:401,415). For every other SSM layer the chain is a dead branch: never evaluated, never detached, one link per token, each link pinning one scalar buffer. Rate = (#SSM layers − 1) per decode token. The mlx allocator counts objects, not bytes, against the resource limit (mlxbackend/metal/allocator.cpp:142).Three confirming experiments (tiny random model, count-probe build of mlx):
mx.eval(cache.left_padding)each step → count flattens completely (mechanism pinned);Why the server path and not plain
generate():advance()no-ops when both fields areNone— the plainmake_prompt_cache()path. The batch pipeline arms the leak even at batch size 1:_merge_caches(generate.py:1036) →ArraysCache.mergesetsleft_padding = mx.array([0] * B)for all-empty caches (cache.py:709), andfinalize()(which nulls the fields) only runsif max_padding > 0(generate.py:1166) — never for a single request.Why vanilla transformers are immune:
BatchKVCache.offsetis the same shape of per-step mx.array state (cache.py:961), but attention consumes it every step viaself.rope(q, offset=cache.offset), forcing and freeing the chain. A dead chain requires per-step array state the forward pass never reads; only the hybrid-SSMArraysCachebookkeeping qualifies.Fix (#1642)
ArraysCachetracks the cumulativeadvance()decrement as a Python int per metadata field, folded in place into the stored array when that field is read (property access ormake_mask) — metadata-side subtraction wrapped to the metadata dtype's modular range, so integer mask arithmetic is identical to the eager decrement (floating metadata rounds at coalesced, int64-gate-sized chunks rather than per step — not a dtype this pipeline produces);left_padding/lengthsstay mx.array-valued properties.advance()creates zero graph nodes. Validation (measured with the initial build of theadvance()fix; the PR adds review-driven hardening for slower leak variants in the same pattern class, covered by unit tests): the 122B crash config survives with 2,725 objects flat; 205k-token tiny-model endurance flat (+0.007/tok, KV-block noise); greedy decode token-identical to stock (64/64).How it was measured (ml-explore/mlx#3942)
The numbers above come from three read-only counters added to the mlx allocator —
mx.get_active_buffer_count()/get_cache_buffer_count()/get_buffer_histogram()— the probe #1185 asked for. The histogram pinned this leak to the 4-byte size class in a single run (3,268 → 480,470 objects over 10k tokens, every other class flat).Likely duplicates / related issues
PoolingCache/RotatingKVCachepath rather thanArraysCache— if the mechanism is the same, it is a sibling instance of the pattern at a different site (per-step lazy arithmetic on state the forward pass never reads), which this patch does not touch. The probe API makes checking it a one-run exercise.[metal::malloc] Resource limit (499000) exceededon qwen3_5 — descriptor-count leak (not OOM) #1185 — same crash signature in LoRA training on qwen3_5. Training doesn't go throughBatchGenerator, so the trigger there may likewise be a different site with the same pattern; the probe PR is the tool for confirming.mlx_lm.serverdistributed, same error string, different mechanism (MLA).Triage notes for that family:
Resource limit (N) exceeded(mlxallocator.cpp:142): the process's own live-object counter against a limit read once at startup (iogpu.rsrc_limitsysctl, hardcoded 499000 fallback). A constant N across crash reports carries no device-state information. Driver-side refusals surface differently (Unable to allocate, command-buffer errors, SIGABRT in residency-set commit).clear_cache, wired limit) measurably do nothing for this leak — it is an object-count mechanism.Environment
M5 Max 128 GB, macOS 26.5.2 (Darwin 25.5.0), mlx 0.32.0 (pip), mlx-lm 0.31.3 (pip), Python 3.12.
repro_leak.py(self-contained, ~380 lines)