[MLX] Add options for clearing context

[metascroy]

Adds two per-model MLX backend runtime options to bound and reduce memory during long/multi-session serving. clear_cache_interval makes MLXBackend::execute() call mlx::core::clear_cache() every N forward calls (counter on the handle, guarded by the global mutex so no atomic is needed), and skip_mutable_buffer_init lets a multi-session load skip allocating the handle's default mutable-buffer (KV/recurrent state) copy, which is dead weight once mlx_mutable_state.h allocates per-session buffers — eliminating one full KV-cache copy. Both are delivered per-delegate via the existing LoadBackendOptionsMap/get_runtime_spec path and exposed through a shared backends/mlx/runtime/backend_options.h. Skipping is gated for safety: a new auto-generated for_each_tid walker (added to generate.py + the MLXLoader templates, covering all op nodes) powers init_chain_references_mutable_buffer(), and init() errors out cleanly if skip_mutable_buffer_init is set while the program's init chain references mutable buffers. mutable_state_bytes_per_session is rebased onto program metadata (with a per-tensor size clamp to avoid overflow from malformed payloads) so capacity accounting stays correct when default buffers are skipped. Finally, the gemma4_31b runner enables these in build_gemma_module: clear_cache_interval is hardcoded to 1, and skip_mutable_buffer_init is set automatically when max_sessions > 1.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20609

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 87 Pending

As of commit 0570950 with merge base 0f3303f ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 3m 40s —— View job

---

Code Review: MLX context-clearing options

I reviewed all 7 changed files. This is a well-constructed PR — the safety gating around skip_mutable_buffer_init, the overflow clamping, and the find()-not-operator[] invariant checks show real care. Findings below, ordered by importance. Nothing here is a hard blocker.

Correctness / robustness

1. skip_mutable_buffer_init is silently unsafe if set without a multi-session owner — MLXBackend.cpp:277-294

The backend trusts the caller to only set this option when a multi-session owner (mlx_mutable_state) is active for the load. If skip_mutable_buffer_init=true but no load scope is active, then:

• handle->mutable_buffers stays empty,
• mutable_state_note_handle() is a no-op (tl_loading_ctx == kInvalidMutableContext), so the handle is never registered,
• on execute(), mutable_state_rebind_for_execute() finds no context and returns Error::Ok keeping the empty default buffers (mlx_mutable_state.cpp:286),
• the program then runs against empty mutable buffers → crash or silent garbage.

In this PR it's correctly gated (the gemma runner only sets it under max_sessions > 1, the same condition that creates the owner), so it's safe as wired up today. But the backend can't detect the misuse, and the two decisions live in different functions. Consider either a defensive check (e.g. error if skip was requested but mutable_state_note_handle didn't register the handle / no owner is active), or at minimum a comment cross-referencing the coupling. The doc comment in backend_options.h:28-32 helps but the runtime can't enforce it.

2. Two independent config.max_sessions > 1 checks must stay in sync — gemma4_31b_engine.cpp:156 and :723

create() decides whether to build the mutable_state owner, and build_gemma_module() independently decides whether to set skip_mutable_buffer_init. Both gate on max_sessions > 1, and build_gemma_module runs inside with_load_scope only when the owner exists — so they're consistent. But because the skip flag is only safe when the load scope is active (see #1), these two conditions are an invariant that isn't expressed in one place. If one is ever changed without the other, you get the bug in #1. A short comment tying them together, or threading a single bool multi_session into build_gemma_module, would make this robust.

Looks correct (verified)

• clear_cache cadence (MLXBackend.cpp:475-480): the second lock_guard is a fresh, non-nested scope (first scope closes at :461), so no deadlock. Counter is non-atomic but only touched under the mutex — correct. ++execute_count_ % clear_cache_interval_ mixes uint64_t/int but interval_ > 0 is guaranteed, so the promotion is well-defined. clear_cache() runs only after this handle's outputs are wait()-ed and state.reset() is done, and it reclaims only free cached buffers, so other handles' in-flight GPU work is unaffected.
• for_each_tid generation (generate.py:1016-1049): covers all Tid-bearing kinds (tid, optional_tid, list_tid, vid_or_tid, int_or_vid_or_tid); NoopNode correctly emits an empty case (no default fallthrough). Because it's auto-generated from the union, a new op without regen would hit the default: throw, which is caught in init() → Error::InvalidProgram. Good failure mode.
• init_chain_references_mutable_buffer (MLXLoader.h.tmpl:224-273): correctly follows nested SCAN/IF chains (which for_each_tid intentionally doesn't), with visited[] cycle protection and bounds checks on chain indices. The IfNode uint32 indices are bounds-checked by chain_idx >= num_chains. Early-out on num_mutable_buffer_tensors == 0 is correct.
• mutable_state_bytes_per_session (mlx_mutable_state.cpp:132-188): the per-tensor clamp to kMaxAllocationBytes correctly prevents accumulator overflow — once clamped, subsequent dims keep it clamped, and d == 0/dynamic dims are handled. Summing across all handles in a context is correct accounting (each handle gets its own per-session buffer set); for the gemma MLX path it's a single forward handle anyway.

Minor

• No tests. The new for_each_tid walker and init_chain_references_mutable_buffer are pure, easily unit-testable functions guarding a memory-safety decision. Given the footgun in Add support for quantized LeakyReLU #1, a small test (init chain that does/doesn't reference a mutable buffer → expected skip/error) would be worthwhile. Were the SCAN/IF recursion paths exercised manually?

Fix item #1 →

---

Summary: solid PR. Main suggestion is to make the skip_mutable_buffer_init ↔ multi-session-owner coupling defensively enforced or clearly documented in one place, and to add a small test for the init-chain walker.
• branch mlx-clear-cxt

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[SS-JIA]

Here's a comment by claude; not sure if valid or not 😅

---

On clear_cache_interval — consider set_cache_limit and the per-handle vs. global granularity

The periodic clear_cache() works, but I think it's worth reconsidering both the mechanism and where the knob lives, because mlx::core::clear_cache() (and set_cache_limit) operate on the process-global singleton allocator, while this option is keyed per-model/per-handle.

Two consequences of that mismatch:

1. The cadence isn't really "every 256 forwards." Each handle carries its own execute_count_ / clear_cache_interval_, but they all flush the same global pool. With prefill + decode methods (and multiple sessions), the effective flush rate is the sum of all handles' executes, and one handle's flush dumps buffers another handle is about to reuse. The global mutex makes this safe, but the semantics are a bit incoherent — a per-handle counter driving a global side effect.
2. It bounds resident-average, not peak. Between clears the cache can still grow to MLX's default ceiling (≈ the whole memory limit), and each flush is followed by a cold-cache re-allocation spike.

set_cache_limit(small) is arguably a better fit here: it's a global ceiling enforced lazily via LRU eviction on the next malloc, so it (a) is correctly a process-global setting, (b) bounds peak cache footprint continuously, (c) targets exactly the stale distinct-size buffers that accumulate from varying prefill/decode shapes, and (d) avoids the periodic flush/re-alloc spike. The tradeoff is it needs a tuned value (too small → lost reuse / prefill↔decode thrash).

Suggestions, in rough order of preference:

• Switch to a set_cache_limit-style cap (could reuse the same runtime-spec plumbing, just a different key/units), or
• If you keep clear_cache(), apply the cadence once globally rather than per-handle (e.g. drop the per-handle counter, or document that the real cadence is aggregate), and note in the comment that peak between clears is still the default ceiling.

Either way, a one-line note in backend_options.h that these knobs affectt just this delegate handle) would save the next reader some confusion.

[metascroy]

Good point on the per-handle/global mismatch; I've kept clear_cache() (it mirrors mlx-lm's own generate_step, which calls mx.clear_cache() every 256 tokens, mlx_lm/generate.py:467) and added comments in backend_options.h and the execute() flush site clarifying that the flush is process-global while the counter is per-handle (exact cadence with one handle, aggregate with several) and that it bounds resident-average, with set_wired_limit/cache-limit noted as a possible future peak cap.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 41s —— View job

---

Code Review: MLX context-clearing options (round 2)

Re-reviewed the current state of the branch against origin/main. Both findings from my earlier review are now addressed — the defensive gating (mutable_state_load_scope_active()) and the granularity comments in backend_options.h. New code below; ordered by importance. One item is worth confirming before merge.

Worth confirming

1. kMLXClearCacheInterval = 1 clears the cache on every forward — is that intended? — gemma4_31b_engine.cpp:147

constexpr int kMLXClearCacheInterval = 1;

The PR description says "clear_cache_interval is hardcoded to 256," and your own review reply cites mlx-lm's generate_step calling mx.clear_cache() every 256 tokens. But the committed value is 1, so (++execute_count_ % 1) == 0 is true on every execute and clear_cache() runs after every forward. That fully defeats inter-step buffer reuse — every decode step pays a cold-cache realloc spike, which is exactly the worst case the backend_options.h comment warns about ("each flush is followed by a cold-cache realloc"). This reads like a debugging value left in. If 256 was intended, please bump it. Fix this →

Looks correct (verified)

• Defensive gating for skip_mutable_buffer_init (MLXBackend.cpp:286-303): the new mutable_state_load_scope_active() check (mlx_mutable_state.cpp:257) rejects the flag when no owner is active for the load, and the init_chain_references_mutable_buffer() check rejects it when the init chain touches mutable buffers. Both throw → caught → Error::InvalidProgram. This closes the footgun from round 1 cleanly. The runner now threads a single multi_session bool (gemma4_31b_engine.cpp:732) as the source of truth for both the owner and the skip flag, so they can't diverge — exactly the fix suggested.

• for_each_tid generation (generate.py:1009-1041): the kind == 2 discriminator for int_or_vid_or_tid matches the codegen (from_tid → kind=2, verified at generate.py:366); vid_or_tid uses !is_vid. All Tid-bearing kinds covered (tid, optional_tid, list_tid, vid_or_tid, int_or_vid_or_tid); non-Tid kinds emit nothing. Empty-body cases (e.g. NoopNode) correctly skip the std::get<> line, and the default: throw is the right failure mode for an un-regenerated new op (caught in init()).

• init_chain_references_mutable_buffer (MLXLoader.h.tmpl:222-273): Tid range [output_end, mutable_buffer_end) matches the documented layout; visited[] cycle guard, chain_idx >= num_chains bounds check, and the early-out on num_mutable_buffer_tensors == 0 are all correct. Worklist follows nested SCAN/IF chains that for_each_tid deliberately doesn't — the IfNode uint32 indices are bounds-checked at pop time, so pushing them unconditionally is safe.

• clear_cache cadence (MLXBackend.cpp:485-491): fresh non-nested lock_guard scope (no deadlock), counter only touched under the mutex (non-atomic is fine), runs after state.reset(). Correct as written — see item Add support for quantized LeakyReLU #1 only for the value, not the mechanism.

• mutable_state_bytes_per_session (mlx_mutable_state.cpp:136-188): the rebase onto program metadata is the right call now that default buffers can be skipped (the old code read default_buffers, which is null under skip). Per-tensor clamp to kMaxAllocationBytes prevents accumulator overflow, d == 0 and dynamic dims handled.

Minor

• Still no tests. for_each_tid and init_chain_references_mutable_buffer are pure functions guarding a memory-safety decision, and the SCAN/IF recursion is non-trivial. A small unit test (init chain that does / doesn't reference a mutable buffer → expected error / skip) would lock in the safety contract. Were the SCAN/IF paths exercised manually?

---

Summary: the round-1 concerns are resolved well. The only thing I'd block on is confirming kMLXClearCacheInterval = 1 vs. the intended 256 — as committed it disables cache reuse entirely.
• branch mlx-clear-cxt