MLX: bound the lazy graph so a method's peak memory is not its whole graph

.github/workflows/mlx.yml

@@ -60,7 +60,7 @@ jobs:
         ${CONDA_RUN} python install_executorch.py > /dev/null
         # The sanitizers fail on github VM runner, but pass on real device
         # TODO: figure out why
-        ${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
+        ${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
         echo "::endgroup::"
 
         ${CONDA_RUN} pip list
@@ -70,7 +70,8 @@ jobs:
         echo "::endgroup::"
 
         echo "::group::Build test runners"
-        ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_metallib_path_test mlx_mutable_state_test mlx_sequence_cache_test mlx_batched_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 ))
+        # Build all configured targets so newly registered tests are built too.
+        ${CONDA_RUN} cmake --build cmake-out -j$(( $(sysctl -n hw.ncpu) - 1 ))
         echo "::endgroup::"
 
         echo "::group::Check MLX artifact sizes"
@@ -104,18 +105,12 @@ jobs:
           exit 1
         fi
 
-        echo "::group::Run SwiftPM metallib path unit test"
-        ./cmake-out/backends/mlx/test/mlx_metallib_path_test
-        echo "::endgroup::"
-
-        echo "::group::Run mutable-state (multi-session) unit test"
-        ./cmake-out/backends/mlx/test/mlx_mutable_state_test
-        echo "::endgroup::"
-
-        echo "::group::Run off-graph KV-cache op test"
-        ./cmake-out/backends/mlx/test/mlx_sequence_cache_test
-        ./cmake-out/backends/mlx/test/mlx_batched_sequence_cache_test
-        ./cmake-out/backends/mlx/test/mlx_cell_cache_test
+        echo "::group::Run registered MLX C++ tests"
+        ${CONDA_RUN} python backends/mlx/test/export_multi_thread_test_model.py /tmp/multi_thread_test_model.pte
+        ET_TESTING_MODEL_PATH=/tmp/multi_thread_test_model.pte \
+          ET_TESTING_NUM_THREADS=50 \
+          ET_PREDICTIONS_PER_THREAD=100 \
+          ${CONDA_RUN} ctest --test-dir cmake-out/backends/mlx/test --output-on-failure --no-tests=error
         echo "::endgroup::"
 
         echo "::group::Run op unit tests"
@@ -135,14 +130,6 @@ jobs:
           -v
         echo "::endgroup::"
 
-        echo "::group::Run multi-thread stress test"
-        ${CONDA_RUN} python backends/mlx/test/export_multi_thread_test_model.py /tmp/multi_thread_test_model.pte
-        ET_TESTING_MODEL_PATH=/tmp/multi_thread_test_model.pte \
-          ET_TESTING_NUM_THREADS=50 \
-          ET_PREDICTIONS_PER_THREAD=100 \
-          ./cmake-out/backends/mlx/test/multi_thread_test_runner
-        echo "::endgroup::"
-
         echo "::group::Run custom_kernel_ops op tests"
         # Run every custom_kernel_ops/**/test/test_*.py via its OpTestCase `run`
         # CLI. Recurses into per-format subpackages (e.g. gguf/test), so adding a

backends/mlx/runtime/MLXBackend.cpp

@@ -286,6 +286,23 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
         handle->clear_cache_interval_ = spec.get();
       }
 
+      // Per-model lazy-graph evaluation threshold (optional runtime spec,
+      // keyed per delegate). Configured here, before the init chain runs
+      // below, so the init chain is covered by the same setting. 0/unset
+      // disables the mechanism and is the default.
+      if (auto spec = context.get_runtime_spec<int>(kEvalThresholdBytesKey);
+          spec.ok()) {
+        const int bytes = spec.get();
+        if (bytes < 0) {
+          throw std::runtime_error(
+              std::string(kEvalThresholdBytesKey) +
+              " must be >= 0 (0 disables the mechanism), got " +
+              std::to_string(bytes));
+        }
+        handle->interpreter.set_eval_threshold_bytes(
+            static_cast<size_t>(bytes));
+      }
+
       if (!processed || !processed->data() || processed->size() == 0) {
         throw std::runtime_error("init: null or empty delegate payload");
       }

backends/mlx/runtime/MLXInterpreter.h

@@ -19,6 +19,11 @@
 #include <mlx/mlx.h>
 #include <mlx/ops.h>
 
+#include <algorithm>
+#include <atomic>
+#include <cstdint>
+#include <limits>
+
 namespace executorch {
 namespace backends {
 namespace mlx {
@@ -1926,60 +1931,180 @@ inline void exec_argpartition(
 
 class Interpreter {
  public:
+  // Threshold in bytes of pending intermediates before the live per-execution
+  // tensors are evaluated. 0 (the default) disables the mechanism entirely:
+  // no traversal, no nbytes() queries, no accumulation, no root collection.
+  // Set from the kEvalThresholdBytesKey runtime spec, before init() runs the
+  // init chain. See backend_options.h for why this is a threshold and not a
+  // hard limit.
+  void set_eval_threshold_bytes(size_t bytes) {
+    eval_threshold_bytes_ = bytes;
+  }
+  size_t eval_threshold_bytes() const {
+    return eval_threshold_bytes_;
+  }
+
+  // Test-only instrumentation: counts calls to accumulate_instruction_bytes.
+  // Only ever incremented on the enabled path, so the disabled path can be
+  // asserted to be free of the accounting work.
+  static uint64_t accounting_calls() {
+    return accounting_calls_.load(std::memory_order_relaxed);
+  }
+  static void reset_accounting_calls() {
+    accounting_calls_.store(0, std::memory_order_relaxed);
+  }
+
   void run(
       const MLXProgram& prog,
       ExecutionState& st,
       StreamOrDevice stream = {}) const {
     run_chain(prog, prog.main_chain_idx, st, stream);
   }
 
+  // Entry point: owns the pending-bytes counter for this execution.
   void run_chain(
       const MLXProgram& prog,
       uint32_t chain_idx,
       ExecutionState& st,
       StreamOrDevice stream = {}) const {
+    size_t pending_bytes = 0;
+    run_chain(
+        prog,
+        chain_idx,
+        st,
+        stream,
+        pending_bytes,
+        /*accumulate_only=*/false);
+  }
+
+  // Nested chains (IF branches, SCAN bodies) share the caller's counter and
+  // pass accumulate_only=true: they add their instructions' estimates but must
+  // not trigger a threshold evaluation themselves. Deferring the check until
+  // control returns to the enclosing chain means a SCAN's collected outputs
+  // have been stacked into state and are reachable from the evaluation roots;
+  // evaluating inside the body could otherwise reset the shared counter while
+  // earlier outputs were still retained only in `collected`.
+  //
+  // accumulate_only suppresses only the threshold-triggered evaluation here.
+  // Op-internal evaluations are unaffected.
+  void run_chain(
+      const MLXProgram& prog,
+      uint32_t chain_idx,
+      ExecutionState& st,
+      StreamOrDevice stream,
+      size_t& pending_bytes,
+      bool accumulate_only) const {
     if (chain_idx >= prog.instruction_chains.size()) {
       throw std::runtime_error(
           "run_chain: chain_idx " + std::to_string(chain_idx) +
           " out of range (num_chains=" +
           std::to_string(prog.instruction_chains.size()) + ")");
     }
     const auto& chain = prog.instruction_chains[chain_idx];
+    const size_t threshold = eval_threshold_bytes_;
     size_t idx = 0;
     for (const auto& instr : chain) {
       st.begin_op(idx, op_name(instr.op));
       if (instr.op == OpCode::SCAN) {
-        exec_scan(prog, std::get<ScanNode>(instr.node), st, stream);
+        exec_scan(
+            prog, std::get<ScanNode>(instr.node), st, stream, pending_bytes);
       } else if (instr.op == OpCode::IF) {
-        exec_if(prog, std::get<IfNode>(instr.node), st, stream);
+        exec_if(prog, std::get<IfNode>(instr.node), st, stream, pending_bytes);
       } else {
         dispatch(instr, st, stream);
       }
       st.end_op();
       ++idx;
+
+      if (threshold != 0) {
+        // SCAN and IF already accumulated their own child instructions through
+        // the shared counter; charging the parent for them again would double
+        // count.
+        if (instr.op != OpCode::SCAN && instr.op != OpCode::IF) {
+          accumulate_instruction_bytes(instr, st, pending_bytes);
+        }
+        if (!accumulate_only && pending_bytes >= threshold) {
+          evaluate_state_tensors(st);
+          pending_bytes = 0;
+        }
+      }
     }
   }
 
  private:
+  // Charge `pending_bytes` for one instruction. The estimate is the largest
+  // per-execution tensor the instruction touches, which tracks the size of what
+  // it just produced without needing to know which tid is its output.
+  // Constants and mutable buffers are excluded: they are not intermediates and
+  // evaluating does not release them.
+  //
+  // Only ever called when the mechanism is enabled.
+  static void accumulate_instruction_bytes(
+      const Instruction& instr,
+      const ExecutionState& st,
+      size_t& pending_bytes) {
+    accounting_calls_.fetch_add(1, std::memory_order_relaxed);
+    size_t widest = 0;
+    for_each_tid(instr, [&](Tid id) {
+      if (id.idx >= st.num_constants && !st.is_mutable_buffer(id)) {
+        uint32_t slot = st.tensor_index(id);
+        if (slot < st.tensors.size() && st.tensors[slot].has_value()) {
+          widest = std::max(widest, st.tensors[slot]->nbytes());
+        }
+      }
+    });
+    add_saturating(pending_bytes, widest);
+  }
+
+  // Saturating add: a pathological program must not wrap the counter back
+  // under the threshold and silently disable the mechanism.
+  static void add_saturating(size_t& acc, size_t add) {
+    if (add > std::numeric_limits<size_t>::max() - acc) {
+      acc = std::numeric_limits<size_t>::max();
+    } else {
+      acc += add;
+    }
+  }
+
+  // Materialize every live per-execution tensor, releasing the graph that
+  // produced them. Results are unchanged by evaluating early.
+  static void evaluate_state_tensors(ExecutionState& st) {
+    std::vector<::mlx::core::array> live;
+    live.reserve(st.tensors.size());
+    for (auto& t : st.tensors) {
+      if (t.has_value()) {
+        live.push_back(*t);
+      }
+    }
+    if (!live.empty()) {
+      ::mlx::core::eval(live);
+    }
+  }
+
+  size_t eval_threshold_bytes_{0};
+  inline static std::atomic<uint64_t> accounting_calls_{0};
+
   void exec_if(
       const MLXProgram& prog,
       const IfNode& n,
       ExecutionState& st,
-      StreamOrDevice s) const {
+      StreamOrDevice s,
+      size_t& pending_bytes) const {
     // Select one branch at runtime based on the integer condition.
     // Nonzero -> then_chain, zero -> else_chain. The selected chain's
     // instructions write the output slot(s) directly.
     const int64_t cond = resolve_int(n.cond, st);
     const uint32_t chain_idx =
         (cond != 0) ? n.then_chain_idx : n.else_chain_idx;
-    run_chain(prog, chain_idx, st, s);
+    run_chain(prog, chain_idx, st, s, pending_bytes, /*accumulate_only=*/true);
   }
 
   void exec_scan(
       const MLXProgram& prog,
       const ScanNode& n,
       ExecutionState& st,
-      StreamOrDevice s) const {
+      StreamOrDevice s,
+      size_t& pending_bytes) const {
     int axis = n.scan_axis;
     int T_int = st.const_tensor_ref(n.originals[0]).shape(axis);
     size_t T = static_cast<size_t>(T_int);
@@ -2001,7 +2126,13 @@ class Interpreter {
                 s));
       }
 
-      run_chain(prog, static_cast<uint32_t>(n.body_chain_idx), st, s);
+      run_chain(
+          prog,
+          static_cast<uint32_t>(n.body_chain_idx),
+          st,
+          s,
+          pending_bytes,
+          /*accumulate_only=*/true);
 
       for (size_t i = 0; i < num_outputs; ++i) {
         collected[i].push_back(st.const_tensor_ref(n.outputs[i]));
@@ -2011,6 +2142,21 @@ class Interpreter {
     for (size_t i = 0; i < num_outputs; ++i) {
       st.set_tensor(n.outputs[i], ::mlx::core::stack(collected[i], axis, s));
     }
+
+    // The stacked outputs are new allocations the body's per-instruction
+    // estimates never saw, so charge for them here. Guarded like every other
+    // piece of the accounting.
+    if (eval_threshold_bytes_ != 0) {
+      for (size_t i = 0; i < num_outputs; ++i) {
+        const Tid id = n.outputs[i];
+        if (id.idx >= st.num_constants && !st.is_mutable_buffer(id)) {
+          uint32_t slot = st.tensor_index(id);
+          if (slot < st.tensors.size() && st.tensors[slot].has_value()) {
+            add_saturating(pending_bytes, st.tensors[slot]->nbytes());
+          }
+        }
+      }
+    }
   }
   void dispatch(const Instruction& instr, ExecutionState& st, StreamOrDevice s)
       const {

backends/mlx/runtime/backend_options.h

@@ -42,6 +42,36 @@ inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval";
 // errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle.
 inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init";
 
+// Per-model runtime-spec key. Value N means: while running a method, evaluate
+// the live per-execution tensors once the intermediates produced since the last
+// evaluation exceed N bytes. 0/unset disables the mechanism entirely and is the
+// default.
+//
+// WHY: MLX is lazy. Interpreter::dispatch only builds graph nodes, and nothing
+// is materialized until MLXBackend::execute calls async_eval on the method
+// outputs, so for a long instruction chain every intermediate in the method is
+// live at the same instant. Whisper-small's 495-instruction encode peaks at
+// 1105 MB of MLX allocation against 95 MB of steady-state active memory, which
+// is what makes the model unusable on an iPhone (pytorch/executorch#22513).
+// Each evaluation costs a GPU sync, so the cost tracks the NUMBER of
+// evaluations; budgeting bytes rather than counting instructions puts them only
+// in the methods that actually allocate.
+//
+// NOTE that this is a THRESHOLD, not a hard memory limit. It is best-effort
+// evaluation scheduling, and peak footprint can exceed it:
+//   - A long SCAN or IF branch accumulates across its whole body and is only
+//     checked once control returns to the enclosing chain, so it can overshoot
+//     by the size of that body.
+//   - The per-instruction estimate is the largest tensor the instruction
+//     touches, which can overcount (an op that only reads a large tensor is
+//     charged for it) and so can trigger evaluation earlier than the true
+//     pending bytes warrant.
+//   - Ops that evaluate internally reduce the real pending work without
+//     reducing the running estimate.
+// Treat it as a knob to trade GPU syncs against peak memory, and tune it
+// against measurements rather than expecting the value to bound RSS.
+inline constexpr char kEvalThresholdBytesKey[] = "eval_threshold_bytes";
+
 } // namespace mlx
 } // namespace backends
 } // namespace executorch

backends/mlx/test/CMakeLists.txt

@@ -115,6 +115,30 @@ if(EXECUTORCH_MLX_ENABLE_SANITIZERS)
 endif()
 add_test(NAME mlx_mutable_state COMMAND mlx_mutable_state_test)
 
+# Lazy-graph evaluation threshold (kEvalThresholdBytesKey). Drives the
+# interpreter over hand-built programs; no model/tokenizer needed.
+et_cxx_test(
+  mlx_eval_threshold_test
+  SOURCES
+  ${CMAKE_CURRENT_LIST_DIR}/mlx_eval_threshold_test.cpp
+  EXTRA_LIBS
+  mlxdelegate
+  mlx_schema
+  mlx
+)
+target_include_directories(
+  mlx_eval_threshold_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime
+)
+if(EXECUTORCH_MLX_ENABLE_SANITIZERS)
+  target_compile_options(
+    mlx_eval_threshold_test PRIVATE -fsanitize=address,undefined
+                                    -fno-omit-frame-pointer
+  )
+  target_link_options(
+    mlx_eval_threshold_test PRIVATE ${_mlx_sanitizer_link_options}
+  )
+endif()
+
 # Off-graph KV cache op-level test (no model/tokenizer needed). et_cxx_test
 # links GTest + executorch_core and registers the ctest target.
 et_cxx_test(