Bound the MLX buffer cache instead of letting it track the machine

[cboulay]

The problem

MLX caches every freed buffer in a std::multimap keyed by exact byte size, and reuse_from_cache accepts a candidate only within min(2 * size, size + 2 * page_size) of the request. 2 * page_size is 32 KiB, so above that the 2 * size branch never wins and it degenerates to a near-exact match. Measured against a cached 3000x256 f32 buffer on mlx 0.32.1:

request delta	reused?
+8 KiB	yes
+16 KiB	no
+300 KiB (one 300-sample bin)	no

So a graph whose message length varies mints a permanent new size class per length, in a cache whose default limit is the size of the machine (23347 MiB on a 24 GB host).

This is easy to miss because RSS does not show it. Metal buffers are IOKit allocations, not malloc'd pages. On a 30 kHz feature chain over one simulated hour, seeing 40 distinct message lengths, phys_footprint went 421 -> 6015 MiB while RSS fell 384 -> 119 MiB.

Growth is a staircase — one permanent step per never-before-seen shape — so a least-squares slope understates it badly. Compare first-vs-last.

It is driven by shape diversity, not shape size. Same chain, constant message length:

scenario	mlx_cache	footprint
constant 300 samples (1 shape)	32 MiB	~300 MiB
constant 9000 samples (1 big shape)	366 MiB	836 MiB
10 shapes	493 MiB	901 MiB
30 shapes	3980 MiB	4432 MiB
1 h, 40 shapes	5732 MiB	6015 MiB

The fix: AsArraySettings.mlx_cache_limit_mb (default 512.0)

Applied only when converting to MLX.

A limit rather than periodic clear_cache() because eviction is LRU: the steady-state shape is re-touched every message and stays at the head of the list, while one-off shapes from a stall fall to the tail and are freed first. So the limit evicts precisely the size classes worth losing and keeps the one worth keeping. Verified — after 39 rare shapes blow past a 128 MiB limit, the hot 300-sample allocation runs at 0.97x its former speed, i.e. unaffected.

20 min stream, stalls every 300 messages, 256 ch:

config	footprint	cache	throughput
baseline	6206 MiB	5743	9.8x
512 MiB	966	519	11.1x
256 MiB	681	257	11.1x
128 MiB	570	135	11.0x
0	396	0	5.9x (-40%)
clear_cache() ~60 s	2206	1747	10.0x
clear_cache() ~6 s	638	298	10.8x
clear_cache() ~1 s	544	0	10.2x

128-512 MiB is faster than unbounded (less memory pressure), so the limit costs nothing. Only 0 hurts. Periodic clearing is strictly dominated on both axes and needs a policy besides.

Two implementation notes:

• Set in _process, not the Unit's initialize(). The transformer is also used bare (offline chains, benchmarks), and set_cache_limit does not survive a spawn — verified: a spawned child reports the 23347 MiB default even when the parent set 256 MiB. It has to run in whichever process actually converts, not the one that built the graph.
• The limit is process-global, so a module-level guard makes it idempotent and warns when a second node asks for a different value, rather than letting whichever unit converted first silently decide.

The smaller half: chunked_scan padded concat

chunk_sizes bounds the set of Metal kernel specializations; it does nothing for buffer size classes, because the scan then undoes it with y_chunk[:, :valid] per chunk and mx.concatenate(...) at exact n_samples. Now it concatenates the padded chunks and trims once, so intermediates land on the multiple-of-chunk_size grid.

Worth 31% in isolation (EWMA core, 400 messages, 40 distinct lengths: 503 -> 349 MiB). Worth about 8% chain-wide (6206 -> 5683 MiB), because most of the cost is message-shaped arrays in stages that never call chunked_scan. Keeping it because it is free and correct, not because it is load-bearing.

Trimming once is sound only because padding is confined to the final chunk: every earlier iteration has remaining > chunk_size, so valid == chunk_size and the chunk is emitted whole. That is an implicit property of the size-selection rule, so the loop now asserts it rather than trusting the next edit to preserve it.

Combined

Full chain, 20 min, stalls every 300 messages: 6206 -> 938 MiB, a 6.6x reduction. The limit is doing nearly all of it.

What this does not do

It bounds the cache; it does not reduce the number of shapes. Bounding that means capping message length at the source (e.g. a max_batches on Window), which is deliberately not in this PR — it would hide the growth from pipelines that never window at all.

A padded-message convention across a whole MLX segment (pad at the conversion in, carry valid_length, trim at the conversion out) would be the real fix and measured -79%, but it was rejected as unsafe: a new edge added to a live graph could bypass the trimming conversion and silently process padding as data.

Testing

4096 passed, 6 skipped. New: tests/unit/test_mlx_metal_common.py (exact-length contract across size sets and lengths, the pad-only-on-tail invariant, and a cache comparison against an inline copy of the pre-change implementation with identical held-live inputs — comparing against varying inputs instead measures input diversity, which this change does not address, and fails at 17.7x); plus AsArray cases in tests/unit/test_asarray.py covering the default, the numpy-target no-op, idempotency, the conflict warning, and that the cache is genuinely capped under size churn.

Judgment call worth a second opinion

The default is 512.0, not None. A library silently setting a process-global on first use can surprise, but a None default only helps people who already know the knob exists — and nobody does, because the growth is invisible in RSS. The accepted risk: large-batch offline MLX work sharing the process gets a 512 MiB cap where it had unlimited, which could thrash. Called out in the docstring. Trivial to flip to opt-in.

All measurements are from an in-process synthetic harness on an M-series host (no ezmsg graph, no SHM), so the throughput column reflects allocator cost rather than live scheduling.