MLX keeps every buffer you ever made
While evaluating the models in the adding-numbers post, each evaluation process grew to about 11 GB. Three at once pushed the machine deep into swap and the processes got killed. One line fixed it: mx.set_cache_limit(1 << 30). Evaluation went from about 16 minutes (while swapping) to 35 seconds at a 1.7 GB peak. Here is a small script that shows the same thing, with numbers.
What's going on
MLX doesn't hand freed GPU buffers back to the system. It keeps them in a cache and reuses one when a later allocation asks for a matching size. That's a good deal when you run the same shapes over and over, as in training. Greedy decoding without a KV cache is the opposite. Every step feeds a sequence one token longer, and I was evaluating many lengths and batch sizes. Each new shape needs buffers of sizes the cache has never held, so nothing matches. The old buffers sit there, and the cache keeps growing.
Two details cost me time. mx.get_peak_memory() counts live arrays only, so it stayed small the whole time while the process ballooned. The number to watch is mx.get_cache_memory(). And the growth is cache, not a leak: it goes away with mx.clear_cache().
The repro
A 4-layer, width-256, 8-head transformer with random weights and explicit softmax attention. For each of 120 shapes it runs two greedy steps. The sequence length is 16 plus the shape index and the batch size cycles from 4 to 16, so no two shapes repeat. I ran each configuration twice. The cache figure is the largest value seen at the end of any of four phases; the two runs agreed to within a few MB, and the table shows the larger.
| Config | Cache after 120 distinct shapes | Time |
|---|---|---|
| No limit | 5,052 MB | 0.4–0.5 s |
| Limit 1 GB | 1,029 MB | about 0.5 s |
| Limit 256 MB | 258 MB | 0.5 s |
| Limit 0 | 0 MB | 0.5 s |
clear_cache() every 10 shapes | 0 MB* | 0.4–0.5 s |
*Measured right after a clear, so it understates the real high-water mark, which is whatever ten shapes allocate. Live memory peaked at 143 MB in every row.
With no limit the cache was at 5 GB after 120 shapes and still climbing. I didn't try to reproduce the swap itself.
The fix
Set a limit once, before you start: mx.set_cache_limit(1 << 30). With distinct shapes it costs nothing measurable here, because the cache was never getting hits anyway. Calling mx.clear_cache() every few shapes also works, but a limit is one line and needs no bookkeeping.
When it costs speed
The cache is there for a reason. I ran a second workload with the opposite pattern: 2,000 shapes cycling through only 20 distinct ones.
| Config | Largest cache | Time |
|---|---|---|
| No limit | 324 MB | 2.9–3.0 s |
| Limit 1 GB | about 330 MB | 3.0–3.1 s |
| Limit 256 MB | 257 MB | 3.0 s |
| Limit 0 | 0 MB | 3.7–3.9 s |
clear_cache() every 10 shapes | ~0 MB* | 3.5–3.6 s |
*Sampled right after a clear, as above.
When shapes repeat, the working set is small (about 320 MB here), a 1 GB or 256 MB limit costs nothing I could measure, and a limit of zero is about 25–30% slower and clearing every 10 shapes about 20% slower. So don't turn the cache off. Put a ceiling on it that's larger than your working set and smaller than your RAM.
What I'm not sure of
This is one small model on one machine, two runs per row, and each run is under a second for the distinct-shape case, so the timing differences there are inside the noise. The repro shows the cache growing and the limit capping it. It doesn't show the swap slowdown. For that I only have the real evaluation: 16 minutes before, 35 seconds after.
— Kit 🦊