Memory leak when compiled function contains multi-output op #3932

Summary

When a function that is compiled with mx.compile contains a multi-output primitive
(mx.split, and by construction any other make_arrays op) and closes over arrays that were
evaluated outside the compiled function, dropping the compiled function and the captured
arrays does not release their memory. mx.get_active_memory() stays elevated permanently and
accumulates without bound across repetitions.

Replacing the single mx.split with two equivalent single-output slices removes the leak
entirely, leaving the numerics bit-identical.

Environment

Minimal reproduction

The script below is the whole report. The fuller ablation harness that produced the tables — same
model, one flag per ingredient — is committed beside this document as
1M3R3-09-T19-repro.py.

import gc
import mlx.core as mx

BLOCKS, DIM, BATCH = 24, 512, 8          # 24 MiB of "weights"
x0 = mx.random.normal((BATCH, DIM)); mx.eval(x0)

def make_forward(weights):
    def forward(x):
        for w in weights:
            h = x @ w
            a, b = mx.split(h, 2, axis=-1)          # <-- multi-output primitive
            x = mx.concatenate([b, a], axis=-1)
        return x
    return forward

for cycle in range(10):
    mx.synchronize(); base = mx.get_active_memory()

    weights = [mx.random.normal((DIM, DIM)) for _ in range(BLOCKS)]
    fwd = make_forward(weights)

    mx.eval(weights)                                # <-- evaluated OUTSIDE the compiled function

    fn = mx.compile(fwd)
    out = fn(x0); mx.eval(out); del out

    del fn, fwd, weights
    gc.collect()

    mx.synchronize()
    print(cycle, (mx.get_active_memory() - base) // (1024 * 1024), "MiB retained")

Expected vs actual

Expected: each iteration prints 0 MiB retained.

Actual: each iteration prints 24 MiB retained — the entire weight set — and the process
footprint grows by 24 MiB per iteration without limit.

Measured over 200 iterations: 200/200 leaked, 24.00 MiB every time, process residual
4800 MiB. Nothing is ever reclaimed; mx.clear_cache() does not help (the buffers were never
freed to the allocator, so they never enter the recycle pool).

All three ingredients are required

n = 30 iterations per configuration, same machine, same script, one flag changed at a time:

The "two slices" row replaces

a, b = mx.split(h, 2, axis=-1)

with

half = h.shape[-1] // 2
a = h[..., :half]
b = h[..., half:]

and nothing else. It is a complete fix in this reproduction.

Note the third ingredient: if the captured arrays are never evaluated outside the compiled
function, they are owned by the compile-cache tape and released with it, so the leak does not
appear. Any real model that materializes its weights before or alongside a compiled call —
loading them, running a warm-up, computing a reference — satisfies it.

Heap evidence

leaks(1) against the parked reproducer (3 iterations, MallocStackLogging=1):

Process 83244: 50 leaks for 8816 total leaked bytes.

STACK OF 1 INSTANCE OF 'ROOT CYCLE: <malloc in mlx::core::array::make_arrays(
    std::vector<mlx::core::SmallVector<int, 10ul>>, std::vector<mlx::core::Dtype> const&,
    std::shared_ptr<mlx::core::Primitive> const&, std::vector<mlx::core::array> const&)>':

    48 (8.53K) ROOT CYCLE: <malloc in mlx::core::array::make_arrays(...) 0xa0b0addd0> [48]
       44 (7.97K) ROOT CYCLE: <std::__shared_ptr_emplace<mlx::core::array::ArrayDesc> 0xa0b823480> [448]
          CYCLE BACK TO <malloc in mlx::core::array::make_arrays(...) 0xa0b0addd0> [48]
          1 (160 bytes) <std::__shared_ptr_emplace<mlx::core::Split> 0xa0a8f6580> [160]
       3 (528 bytes) ROOT CYCLE: <std::__shared_ptr_emplace<mlx::core::array::ArrayDesc> 0xa0b8239c0> [448]

Unreachable, cyclic, rooted at make_arrays, primitive mlx::core::Split.

leaks reports kilobytes rather than megabytes because it only walks malloc zones: the leaked
nodes are the ArrayDesc / Split / sibling-vector control blocks, while the buffers they pin
are Metal VM regions. vmmap on the same process attributes the retained bytes to
IOAccelerator (graphics).

Analysis

array::make_arrays (mlx/array.cpp:42-58) gives every output of a multi-output primitive a
by-value copy of each of its siblings:

for (size_t i = 0; i < outputs.size(); ++i) {
  auto siblings = outputs;
  siblings.erase(siblings.begin() + i);
  outputs[i].set_siblings(std::move(siblings), i);
}

That is a reference cycle by construction: each ArrayDesc holds a shared_ptr to the others.
~ArrayDesc is aware of this and breaks it — but only for nodes it first judges deletable
(mlx/array.cpp:299-316):

ad.inputs.clear();                                    // :299  unconditional
for (auto& [_, a] : input_map) {
  bool is_deletable =
      (a.array_desc_.use_count() <= a.siblings().size() + 1);   // :301-302
  for (auto& s : a.siblings()) { ... }
  if (is_deletable) {
    for_deletion.push_back(std::move(a.array_desc_));
  }
}

and the cycle-breaking siblings.clear() runs only inside the for_deletion drain
(mlx/array.cpp:328-334).

The failure is the combination of those two lines:

1. ad.inputs.clear() at :299 runs unconditionally, so the parent releases its reference to
   the group before deletability is decided.
2. If use_count() at :301 exceeds the threshold — the count is read without synchronization,
   and the copy in input_map plus any transient reference elsewhere is enough — the group is
   skipped, never enters for_deletion, and therefore never reaches siblings.clear().
3. The group is now referenced by nothing but itself, and it still holds its own inputs, which
   is what pins the upstream weights.

A single-output node cannot hit this: with siblings().size() == 0 the threshold is
use_count() <= 1, and a node rejected there is merely deferred — whoever holds the extra
reference will release it later and ~ArrayDesc will run normally. Only a sibling group can be
orphaned in a state where nothing will ever come back for it.

Compilation is required because the compile-cache tape is what keeps the traced group alive past
the forward pass; without it the graph is transient and torn down while its parents still hold it.
(Marked as inference: the precise reference that inflates the count at :301 was not isolated —
what is measured is that the cycle survives, is unreachable, and is rooted at make_arrays.)

Suggested fix

Break the sibling cycle unconditionally rather than only on the deletable path — i.e. clear
siblings for every entry in input_map that the parent has just released, not only for those
that pass the use_count() test. Alternatively, hold siblings as weak_ptr, which removes the
cycle at the source.

Impact

This affects any long-lived process that builds and drops compiled models: the leak is the full
captured parameter set per drop, permanent, and unbounded. In our case (a 210 MiB speech model
behind mx.compile) a repeating build/drop cycle reached 1.18 GiB retained after 100 iterations
and was still climbing linearly.

Confirmation from the C++ side

The same three ingredients, and the same fix, reproduce through mlx-rs on a real 210 MiB model
(MossFormer2_SE). There the race is intermittent rather than deterministic, which makes the
ablations statistically rather than trivially conclusive:

Fisher p = 1.1 × 10⁻⁷² and 4.1 × 10⁻¹² for the two slice ablations. Note the last two rows: it is
materialization of the captured arrays that arms the bug, not the shape of whatever computation
did the materializing — a bare mx.eval is sufficient.

Replacing the multi-output op cost nothing measurable (median iteration 138 ms vs 139 ms) and left
every output bit-identical.

Also, unrelated to the leak but noticed alongside it: a second compiled shape entry for the same
function costs a full extra copy of the captured parameters in live memory (638.6 MiB vs 427.4 MiB
for one shape, same model). That looks like intended tape behaviour rather than a defect, but it is
worth knowing before compiling several shapes per model.

AI Disclosure

This issue was identified by Claude Opus 5 as the root cause of a memory leak in a downstream project that depends on MLX. We have reworked our code to avoid this issue, and I'm submitting the report Claude wrote for your information and any follow-up you might want to take. Thanks for your work on this library!