Assert the bound MLX honours, not the limit it was given

[cboulay]

Follow-up to #236, which was insufficient. Sorry — that one is on me.

What #236 got wrong

It forced a trim with one extra allocation before reading the total. That worked 12/12 locally and still failed one macOS job in four on #235:

AssertionError: assert 35962884 <= ((32 * 1024) * 1024)

The mistake was validating in isolation. Of the six local runs I based it on, only one started above the limit — so I had a sample of one for the path I was fixing, and read six greens as confirmation.

The assertion itself is wrong

MLX evicts down to the limit before admitting a freed buffer, then admits it. The cache is a high-water mark that one buffer can exceed; set_cache_limit never promised a hard ceiling.

Which side of the limit it lands on varies run to run. Measured across twelve fresh processes, the cache after the churn settles at one of two states:

end state
largest churned buffer	17,825,796 B	17.00 MB — under the limit, passes
largest two	35,962,888 B	34.30 MB — over the limit, fails

Neither is a bug. Both are MLX doing what it documents.

The fix

Assert the bound the allocator actually honours — limit + largest_buffer:

CI observed     35962884 B   34.30 MB
old assertion   33554432 B   32.00 MB   -> FAIL
new assertion   51679232 B   49.29 MB   -> pass
headroom        15716348 B   14.99 MB

It's still a meaningful bound, which is the thing worth checking about a looser assertion:

• the same churn with no limit at all leaves ~500 MB cached — ten times this ceiling;
• raising mlx_cache_limit_mb to 100 GB makes the test fail.

So it still proves what it was written to prove.

Honesty about the evidence

I could not reproduce the failure locally, with or without #236's fix — ten isolated runs and three full-suite runs all pass either way. The argument here rests on arithmetic against CI's own failing value, not on a green local run. If it recurs, the next step is to log get_cache_memory() in CI rather than guess again.

Unrelated, found while checking this

tests/unit/test_buffer_recycling.py::test_binned_aggregate_no_bin_completed_does_not_retain fails on clean dev — a second pre-existing red, nothing to do with MLX. It passes on my #233 stack, which fixes a gap in tests/helpers/recycled_shm.py (_detach didn't copy axes). Flagging it here rather than folding it in; it resolves when that stack lands.

🤖 Generated with Claude Code