[MLX] Add options for clearing context (#20609)

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.

Author: metascroy

backends/mlx/runtime/MLXBackend.cpp

@@ -19,6 +19,8 @@
 
 #include <mlx/mlx.h>
 
+#include <executorch/backends/mlx/runtime/backend_options.h>
+
 #include <cstring>
 #include <limits>
 #include <memory>
@@ -169,6 +171,12 @@ struct MLXHandle {
   Interpreter interpreter;
   ::mlx::core::Stream stream; // Dedicated GPU stream for this handle
 
+  // Periodically call mlx::core::clear_cache() every N execute() calls to
+  // release MLX's cached buffer pool. 0 disables. Counter is non-atomic: it is
+  // only touched inside mlx_global_mutex() (see execute()).
+  int clear_cache_interval_{0};
+  uint64_t execute_count_{0};
+
   // Keep the constant buffers alive for zero-copy constants
   // Each FreeableBuffer must outlive the MLX arrays that reference it
   std::vector<FreeableBuffer> constant_buffers;
@@ -211,6 +219,13 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
     try {
       new (handle) MLXHandle();
 
+      // Per-model clear-cache cadence (optional runtime spec, keyed per
+      // delegate)
+      if (auto spec = context.get_runtime_spec<int>(kClearCacheIntervalKey);
+          spec.ok() && spec.get() > 0) {
+        handle->clear_cache_interval_ = spec.get();
+      }
+
       if (!processed || !processed->data() || processed->size() == 0) {
         throw std::runtime_error("init: null or empty delegate payload");
       }
@@ -254,8 +269,42 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       processed->Free();
       processed = nullptr;
 
-      // Load mutable buffers (e.g., KV cache)
-      load_mutable_buffers(handle->program, handle->mutable_buffers);
+      // Load mutable buffers (e.g., KV cache). When skip_mutable_buffer_init is
+      // set (multi-session, where mlx_mutable_state.h allocates per-session
+      // buffers), avoid keeping a dead default copy on the handle. This is only
+      // safe if the init chain does not reference mutable buffers — otherwise
+      // it would run against an empty store, so we error out clearly.
+      bool skip_mutable_buffer_init = false;
+      if (auto spec = context.get_runtime_spec<bool>(kSkipMutableBufferInitKey);
+          spec.ok()) {
+        skip_mutable_buffer_init = spec.get();
+      }
+      // skip_mutable_buffer_init is only safe under a multi-session owner: the
+      // per-session manager allocates the buffers and execute() rebinds to
+      // them. Without an active load scope no handle is registered, no
+      // per-session buffers are ever created, and execute() would run against
+      // empty buffers. Reject the misuse loudly here instead of failing later.
+      if (skip_mutable_buffer_init && !mutable_state_load_scope_active()) {
+        ET_LOG(
+            Error,
+            "skip_mutable_buffer_init set without an active multi-session load "
+            "scope; mutable buffers would never be allocated");
+        throw std::runtime_error(
+            "skip_mutable_buffer_init requires an active multi-session owner");
+      }
+      if (skip_mutable_buffer_init &&
+          init_chain_references_mutable_buffer(handle->program)) {
+        ET_LOG(
+            Error,
+            "skip_mutable_buffer_init set but the init chain references "
+            "mutable buffers; cannot skip default mutable-buffer init");
+        throw std::runtime_error(
+            "skip_mutable_buffer_init incompatible with init chain that "
+            "references mutable buffers");
+      }
+      if (!skip_mutable_buffer_init) {
+        load_mutable_buffers(handle->program, handle->mutable_buffers);
+      }
 
       // Bind execution state (reused across execute() calls)
       handle->state.bind(
@@ -433,6 +482,16 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
 
       h->state.reset(); // Release temp GPU buffers back to MLX cache
 
+      // Periodically release MLX's cached buffer pool. The counter increment
+      // and clear_cache() run under the global mutex so they can't race another
+      // handle's in-flight submission, and the counter needs no atomic.
+      if (h->clear_cache_interval_ > 0) {
+        std::lock_guard<std::mutex> lock(mlx_global_mutex());
+        if ((++h->execute_count_ % h->clear_cache_interval_) == 0) {
+          ::mlx::core::clear_cache();
+        }
+      }
+
       return Error::Ok;
     } catch (const std::exception& e) {
       ET_LOG(Error, "MLX execute failed: %s", e.what());
@@ -455,7 +514,7 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
 
 namespace {
 auto cls = MLXBackend();
-Backend backend{"MLXBackend", &cls};
+Backend backend{kMLXBackendId, &cls};
 static auto success_with_compiler = register_backend(backend);
 } // namespace
 

backends/mlx/runtime/backend_options.h

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Shared option keys for the MLX backend. Included by the backend itself
+// (to read the per-model runtime spec) and by callers/runners (to set it via
+// a LoadBackendOptionsMap), keeping the string literals in one place.
+
+#pragma once
+
+namespace executorch {
+namespace backends {
+namespace mlx {
+
+// Backend id under which the MLX backend registers (see MLXBackend.cpp).
+inline constexpr char kMLXBackendId[] = "MLXBackend";
+
+// Per-model runtime-spec key. Value N means: call mlx::core::clear_cache()
+// every N execute() calls to release MLX's cached buffer pool; 0/unset
+// disables.
+//
+// NOTE on granularity: MLX's buffer cache is a process-global singleton, so the
+// flush is global even though this key is read per delegate handle. The counter
+// is per-handle, so with a single MLX handle (the common case — gemma's
+// prefill/decode share one "forward" method) the cadence is exactly "every N
+// forwards"; if a process loads multiple MLX handles, the effective cadence is
+// the aggregate of their executes and any handle's flush frees the shared pool
+// for all. This bounds resident-*average*: between flushes the cache can still
+// grow to MLX's default ceiling, and each flush is followed by a cold-cache
+// realloc. A future set_cache_limit-style key could complement this by bounding
+// peak footprint continuously.
+inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval";
+
+// Per-model runtime-spec key (bool). When true, the handle does not allocate
+// its own default mutable buffers at init() — per-session buffers are managed
+// by mlx_mutable_state.h instead. Only valid for multi-session loads, and only
+// when the program's init chain does not reference mutable buffers (init()
+// errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle.
+inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init";
+
+} // namespace mlx
+} // namespace backends
+} // namespace executorch

backends/mlx/runtime/mlx_mutable_state.cpp

@@ -136,15 +136,52 @@ int64_t mutable_state_bytes_per_session(MutableStateContext ctx) {
     return 0;
   }
   int64_t total = 0;
+  // Per-session mutable buffers are allocated from program metadata
+  // (mutable_buffer_map + tensor_meta), independent of whether the handle kept
+  // a default copy. Compute the estimate from metadata so it stays correct even
+  // when default mutable-buffer init was skipped (skip_mutable_buffer_init).
   for (const auto& kv : it->second.handles) {
-    const MutableBufferData* bufs = kv.second.default_buffers;
-    if (bufs == nullptr) {
+    const MLXProgram* program = kv.second.program;
+    if (program == nullptr) {
       continue;
     }
-    for (const auto& t : bufs->tensors) {
-      if (t.has_value()) {
-        total += static_cast<int64_t>(t->nbytes());
+    for (const auto& slot : program->mutable_buffer_map) {
+      if (slot.slot_type != SlotType::TensorSlot ||
+          slot.idx >= program->tensor_meta.size() ||
+          !program->tensor_meta[slot.idx].has_value()) {
+        continue;
       }
+      const auto& meta = *program->tensor_meta[slot.idx];
+      // Sum sizes from metadata, clamping each tensor to kMaxAllocationBytes so
+      // malformed (oversized) shapes in an untrusted program can't overflow the
+      // accumulator. Real allocations are independently bounded by
+      // check_allocation_bounded at load_mutable_buffers.
+      uint64_t bytes = static_cast<uint64_t>(
+          ::mlx::core::size_of(resolve_dtype(meta.scalar_type)));
+      bool dynamic = false;
+      for (const auto& dim : meta.shape) {
+        if (dim.value < 0) {
+          dynamic = true;
+          break;
+        }
+        const uint64_t d = static_cast<uint64_t>(dim.value);
+        if (d == 0) {
+          bytes = 0;
+          break;
+        }
+        if (bytes > kMaxAllocationBytes / d) {
+          bytes = kMaxAllocationBytes; // clamp; avoids overflow
+        } else {
+          bytes *= d;
+        }
+      }
+      if (dynamic) {
+        continue;
+      }
+      if (bytes > kMaxAllocationBytes) {
+        bytes = kMaxAllocationBytes;
+      }
+      total += static_cast<int64_t>(bytes);
     }
   }
   return total;
@@ -217,6 +254,10 @@ void mutable_state_note_handle(
   handle_ctx()[handle] = tl_loading_ctx;
 }
 
+bool mutable_state_load_scope_active() {
+  return tl_loading_ctx != kInvalidMutableContext;
+}
+
 void mutable_state_forget_handle(const void* handle) {
   std::lock_guard<std::mutex> g(registry_mutex());
   auto hit = handle_ctx().find(handle);

backends/mlx/runtime/mlx_mutable_state.h

@@ -195,6 +195,12 @@ void mutable_state_note_handle(
 
 void mutable_state_forget_handle(const void* handle);
 
+// True if a multi-session load scope (with_load_scope) is active on the current
+// thread. The MLXBackend uses this to reject skip_mutable_buffer_init when no
+// owner is active for the load — without an owner, no per-session buffers are
+// ever allocated and execute() would run against empty mutable buffers.
+bool mutable_state_load_scope_active();
+
 ::executorch::runtime::Error mutable_state_rebind_for_execute(
     const void* handle,
     ExecutionState& state);

backends/mlx/serialization/MLXLoader.h.tmpl

@@ -118,6 +118,25 @@ struct Instruction {
   }
 };
 
+// =============================================================================
+// for_each_tid - invokes cb(Tid) for every tensor id referenced by an
+// instruction's node (AUTO-GENERATED from schema.fbs).
+//
+// NOTE: nested chain-index fields (ScanNode::body_chain_idx,
+// IfNode::then_chain_idx/else_chain_idx) are NOT followed here; callers that
+// need transitive coverage must recurse into those chains themselves.
+// =============================================================================
+
+template <typename F>
+inline void for_each_tid(const Instruction& instr, F&& cb) {
+  switch (instr.op) {
+{{FOR_EACH_TID_CASES}}
+    default:
+      throw std::runtime_error(
+          std::string("for_each_tid: unhandled op ") + op_name(instr.op));
+  }
+}
+
 // =============================================================================
 // SlotVariant for I/O mapping
 // =============================================================================
@@ -194,6 +213,65 @@ struct MLXProgram {
   }
 };
 
+// =============================================================================
+// init_chain_references_mutable_buffer - returns true if the program's init
+// chain references any mutable-buffer Tid, following nested SCAN/IF chains.
+//
+// Used to decide whether the handle's default mutable-buffer initialization can
+// be safely skipped (e.g. when buffers are managed per-session). If the init
+// chain reads/writes mutable buffers, skipping would be unsafe.
+// =============================================================================
+inline bool init_chain_references_mutable_buffer(const MLXProgram& program) {
+  if (program.init_chain_idx < 0 || program.num_mutable_buffer_tensors == 0) {
+    return false;
+  }
+  // Tid layout: Constant -> Input -> Output -> MutableBuffer -> Temp.
+  const uint32_t output_end = program.num_constant_tensors +
+      program.num_input_tensors + program.num_output_tensors;
+  const uint32_t mutable_buffer_end =
+      output_end + program.num_mutable_buffer_tensors;
+  auto is_mutable_buffer = [&](Tid t) {
+    return t.idx >= output_end && t.idx < mutable_buffer_end;
+  };
+
+  const size_t num_chains = program.instruction_chains.size();
+  std::vector<bool> visited(num_chains, false);
+  std::vector<uint32_t> worklist;
+  worklist.push_back(static_cast<uint32_t>(program.init_chain_idx));
+
+  while (!worklist.empty()) {
+    const uint32_t chain_idx = worklist.back();
+    worklist.pop_back();
+    if (chain_idx >= num_chains || visited[chain_idx]) {
+      continue;
+    }
+    visited[chain_idx] = true;
+    for (const auto& instr : program.instruction_chains[chain_idx]) {
+      bool hit = false;
+      for_each_tid(instr, [&](Tid t) {
+        if (is_mutable_buffer(t)) {
+          hit = true;
+        }
+      });
+      if (hit) {
+        return true;
+      }
+      // Follow nested chains (not covered by for_each_tid).
+      if (instr.op == OpCode::SCAN) {
+        const auto& n = std::get<ScanNode>(instr.node);
+        if (n.body_chain_idx >= 0) {
+          worklist.push_back(static_cast<uint32_t>(n.body_chain_idx));
+        }
+      } else if (instr.op == OpCode::IF) {
+        const auto& n = std::get<IfNode>(instr.node);
+        worklist.push_back(n.then_chain_idx);
+        worklist.push_back(n.else_chain_idx);
+      }
+    }
+  }
+  return false;
+}
+
 // =============================================================================
 // FlatBuffer loading functions
 // =============================================================================

backends/mlx/serialization/generate.py

@@ -996,16 +996,59 @@ def generate_cpp_loader_h(schema: FBSSchema) -> str:
         comma = "," if i < len(op_nodes) - 1 else ""
         variant_lines.append(f"    {table.name}{comma}")
 
+    for_each_tid_lines = []
+    for table in op_nodes:
+        for_each_tid_lines.extend(_generate_for_each_tid_case(table))
+
     # Read template and fill placeholders
     header = "\n".join(_file_header("//")) + "\n//\n"
     tmpl = LOADER_H_TMPL.read_text()
     result = tmpl.replace("{{OP_NODE_STRUCTS}}", "\n".join(struct_lines))
     result = result.replace("{{OPCODE_ENUM_VALUES}}", "\n".join(enum_lines))
     result = result.replace("{{OP_NAME_CASES}}", "\n".join(name_lines))
     result = result.replace("{{NODE_VARIANT_TYPES}}", "\n".join(variant_lines))
+    result = result.replace("{{FOR_EACH_TID_CASES}}", "\n".join(for_each_tid_lines))
     return header + result
 
 
+# for_each_tid emitters: invoke cb(Tid) for each Tid-bearing field. Field kinds
+# that carry no Tid (vid, int_or_vid, scalars, strings, int lists) emit nothing.
+def _emit_cpp_for_each_tid(kind: str, name: str) -> List[str]:
+    """Emit for_each_tid callback lines for a field kind ([] if it has no Tid)."""
+    if kind == "tid":
+        return [f"      cb(n.{name});"]
+    if kind == "optional_tid":
+        return [f"      if (n.{name}.has_value()) cb(*n.{name});"]
+    if kind == "list_tid":
+        return [f"      for (const auto& tid_elem : n.{name}) cb(tid_elem);"]
+    if kind == "vid_or_tid":
+        return [f"      if (!n.{name}.is_vid) cb(n.{name}.tid);"]
+    if kind == "int_or_vid_or_tid":
+        return [f"      if (n.{name}.kind == 2) cb(n.{name}.tid);"]
+    return []
+
+
+def _generate_for_each_tid_case(table: FBSTable) -> List[str]:
+    """Generate a switch case for for_each_tid over one op node."""
+    class_name = table.name
+    op_code = _table_name_to_opcode(class_name)
+
+    body: List[str] = []
+    for fld in table.fields:
+        if fld.name.endswith("_is_set"):
+            continue
+        kind = _get_field_kind(fld, table)
+        body.extend(_emit_cpp_for_each_tid(kind, fld.name))
+
+    lines = [f"    case OpCode::{op_code}: {{"]
+    if body:
+        lines.append(f"      const auto& n = std::get<{class_name}>(instr.node);")
+        lines.extend(body)
+    lines.append("      break;")
+    lines.append("    }")
+    return lines
+
+
 def _is_interned_str(table, field_name) -> bool:
     """Whether a string field should be loaded as an interned shared_ptr.