Skip to content

Commit 067632d

Browse files
MLX: bound the lazy graph so a method's peak memory is not its whole graph (pytorch#22932)
Fixes the root cause of pytorch#22513. ## The problem MLX is lazy. `Interpreter::dispatch` only builds graph nodes; nothing is materialized until `MLXBackend::execute` calls `async_eval` on the method outputs (`MLXBackend.cpp:557`). For a long instruction chain that means **every intermediate in the method is live at the same instant**. Whisper-small `encode` is 495 instructions built before a single byte is evaluated. Measured on macOS with `mlx::core::get_{active,peak,cache}_memory()`: | | whisper-small encode | | --- | --- | | MLX peak | **1105.6 MB** | | MLX active, steady | 94.8 MB | | MLX cache | 262.3 MB (at the existing 256 MB cap) | On iOS that peak lands in the app's footprint, which is what makes the model unusable there. ## The fix Evaluate the live per-execution tensors once the intermediates produced since the last barrier exceed a byte threshold. Off by default; opt in per model with the `eval_threshold_bytes` runtime option. Each barrier costs a GPU sync, so the cost tracks the **number** of barriers, not the interval. Budgeting bytes rather than counting instructions puts barriers only in the methods that actually allocate: whisper-small at 512 MB takes **12** barriers in `encode` and **0** in `decode`. An every-32-instruction rule took 15 and 22, and those 22 bought 50 MB on a method that peaks at 258 MB while costing 21% on device. **iPhone 16, whisper-small int8, full pipeline, medians of interleaved rounds:** | arm | peak MB | peak-loaded | pipeline ms | encode | decode x10 | | --- | --- | --- | --- | --- | --- | | no barrier | 1194.4 | 763.7 | 885.2 | 463.2 | 472.6 | | **eval_threshold_bytes 512 MB** | **692.8** | **261.0** | **831.3** | **420.5** | **410.4** | 1.72x lower peak, 2.9x lower execute-phase footprint, and 1.06x **faster** than the unbounded path. **macOS, medians of 4 interleaved rounds x 10 executions, round 1 discarded:** | model | method | peak off -> 512 MB | speed vs off | | --- | --- | --- | --- | | whisper-small | encode | 1105.6 -> 350.1 MB | 1.05x | | whisper-small | decode | 258.0 -> 258.0 MB | 0.97x | | whisper-tiny | encode | 550.4 -> 343.2 MB | 1.07x | | whisper-tiny | decode | 61.8 -> 61.8 MB | 0.93x | | SmolLM2-135M | forward | 1196.5 -> 496.8 MB | 1.00x | Outputs are **bit-identical** to the unbounded path on whisper tiny/base/small x {`encode`, `decode`} and on two MLX LLMs. The threshold is the `eval_threshold_bytes` runtime option, read per delegate handle alongside `clear_cache_interval` and validated at load. It defaults to **0, which disables the mechanism entirely**: no traversal, no `nbytes()` queries, no accumulation, no evaluation-root collection. The measurements above are at 512 MB. ## What this is not Two other candidates were measured and ruled out: - **Buffer cache.** Already capped by `set_cache_limit(256MB)` at init. Forcing `ET_MLX_CACHE_LIMIT_MB=16` moves cache 262 -> 103 MB and leaves peak bit-identical at 1105.6 MB. - **Slot retention.** `ExecutionState::tensors` holds only **7** slots for those 495 instructions; the AOT side already plans and reuses them. A last-use release built on `for_each_tid` freed 5 slots and moved peak by 0 MB. The arrays are pinned by the pending graph, not the slot table. cc @metascroy --------- Co-authored-by: Scott Roy <scroy@meta.com>
1 parent dd8297a commit 067632d

7 files changed

Lines changed: 562 additions & 28 deletions

File tree

‎.github/workflows/mlx.yml‎

Lines changed: 9 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ jobs:
6060
${CONDA_RUN} python install_executorch.py > /dev/null
6161
# The sanitizers fail on github VM runner, but pass on real device
6262
# TODO: figure out why
63-
${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
63+
${CONDA_RUN} cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON -DEXECUTORCH_MLX_ENABLE_SANITIZERS=OFF
6464
echo "::endgroup::"
6565
6666
${CONDA_RUN} pip list
@@ -70,7 +70,8 @@ jobs:
7070
echo "::endgroup::"
7171
7272
echo "::group::Build test runners"
73-
${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 ))
73+
# Build all configured targets so newly registered tests are built too.
74+
${CONDA_RUN} cmake --build cmake-out -j$(( $(sysctl -n hw.ncpu) - 1 ))
7475
echo "::endgroup::"
7576
7677
echo "::group::Check MLX artifact sizes"
@@ -104,18 +105,12 @@ jobs:
104105
exit 1
105106
fi
106107
107-
echo "::group::Run SwiftPM metallib path unit test"
108-
./cmake-out/backends/mlx/test/mlx_metallib_path_test
109-
echo "::endgroup::"
110-
111-
echo "::group::Run mutable-state (multi-session) unit test"
112-
./cmake-out/backends/mlx/test/mlx_mutable_state_test
113-
echo "::endgroup::"
114-
115-
echo "::group::Run off-graph KV-cache op test"
116-
./cmake-out/backends/mlx/test/mlx_sequence_cache_test
117-
./cmake-out/backends/mlx/test/mlx_batched_sequence_cache_test
118-
./cmake-out/backends/mlx/test/mlx_cell_cache_test
108+
echo "::group::Run registered MLX C++ tests"
109+
${CONDA_RUN} python backends/mlx/test/export_multi_thread_test_model.py /tmp/multi_thread_test_model.pte
110+
ET_TESTING_MODEL_PATH=/tmp/multi_thread_test_model.pte \
111+
ET_TESTING_NUM_THREADS=50 \
112+
ET_PREDICTIONS_PER_THREAD=100 \
113+
${CONDA_RUN} ctest --test-dir cmake-out/backends/mlx/test --output-on-failure --no-tests=error
119114
echo "::endgroup::"
120115
121116
echo "::group::Run op unit tests"
@@ -135,14 +130,6 @@ jobs:
135130
-v
136131
echo "::endgroup::"
137132
138-
echo "::group::Run multi-thread stress test"
139-
${CONDA_RUN} python backends/mlx/test/export_multi_thread_test_model.py /tmp/multi_thread_test_model.pte
140-
ET_TESTING_MODEL_PATH=/tmp/multi_thread_test_model.pte \
141-
ET_TESTING_NUM_THREADS=50 \
142-
ET_PREDICTIONS_PER_THREAD=100 \
143-
./cmake-out/backends/mlx/test/multi_thread_test_runner
144-
echo "::endgroup::"
145-
146133
echo "::group::Run custom_kernel_ops op tests"
147134
# Run every custom_kernel_ops/**/test/test_*.py via its OpTestCase `run`
148135
# CLI. Recurses into per-format subpackages (e.g. gguf/test), so adding a

‎backends/mlx/runtime/MLXBackend.cpp‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,23 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
286286
handle->clear_cache_interval_ = spec.get();
287287
}
288288

289+
// Per-model lazy-graph evaluation threshold (optional runtime spec,
290+
// keyed per delegate). Configured here, before the init chain runs
291+
// below, so the init chain is covered by the same setting. 0/unset
292+
// disables the mechanism and is the default.
293+
if (auto spec = context.get_runtime_spec<int>(kEvalThresholdBytesKey);
294+
spec.ok()) {
295+
const int bytes = spec.get();
296+
if (bytes < 0) {
297+
throw std::runtime_error(
298+
std::string(kEvalThresholdBytesKey) +
299+
" must be >= 0 (0 disables the mechanism), got " +
300+
std::to_string(bytes));
301+
}
302+
handle->interpreter.set_eval_threshold_bytes(
303+
static_cast<size_t>(bytes));
304+
}
305+
289306
if (!processed || !processed->data() || processed->size() == 0) {
290307
throw std::runtime_error("init: null or empty delegate payload");
291308
}

‎backends/mlx/runtime/MLXInterpreter.h‎

Lines changed: 152 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@
1919
#include <mlx/mlx.h>
2020
#include <mlx/ops.h>
2121

22+
#include <algorithm>
23+
#include <atomic>
24+
#include <cstdint>
25+
#include <limits>
26+
2227
namespace executorch {
2328
namespace backends {
2429
namespace mlx {
@@ -1926,60 +1931,180 @@ inline void exec_argpartition(
19261931

19271932
class Interpreter {
19281933
public:
1934+
// Threshold in bytes of pending intermediates before the live per-execution
1935+
// tensors are evaluated. 0 (the default) disables the mechanism entirely:
1936+
// no traversal, no nbytes() queries, no accumulation, no root collection.
1937+
// Set from the kEvalThresholdBytesKey runtime spec, before init() runs the
1938+
// init chain. See backend_options.h for why this is a threshold and not a
1939+
// hard limit.
1940+
void set_eval_threshold_bytes(size_t bytes) {
1941+
eval_threshold_bytes_ = bytes;
1942+
}
1943+
size_t eval_threshold_bytes() const {
1944+
return eval_threshold_bytes_;
1945+
}
1946+
1947+
// Test-only instrumentation: counts calls to accumulate_instruction_bytes.
1948+
// Only ever incremented on the enabled path, so the disabled path can be
1949+
// asserted to be free of the accounting work.
1950+
static uint64_t accounting_calls() {
1951+
return accounting_calls_.load(std::memory_order_relaxed);
1952+
}
1953+
static void reset_accounting_calls() {
1954+
accounting_calls_.store(0, std::memory_order_relaxed);
1955+
}
1956+
19291957
void run(
19301958
const MLXProgram& prog,
19311959
ExecutionState& st,
19321960
StreamOrDevice stream = {}) const {
19331961
run_chain(prog, prog.main_chain_idx, st, stream);
19341962
}
19351963

1964+
// Entry point: owns the pending-bytes counter for this execution.
19361965
void run_chain(
19371966
const MLXProgram& prog,
19381967
uint32_t chain_idx,
19391968
ExecutionState& st,
19401969
StreamOrDevice stream = {}) const {
1970+
size_t pending_bytes = 0;
1971+
run_chain(
1972+
prog,
1973+
chain_idx,
1974+
st,
1975+
stream,
1976+
pending_bytes,
1977+
/*accumulate_only=*/false);
1978+
}
1979+
1980+
// Nested chains (IF branches, SCAN bodies) share the caller's counter and
1981+
// pass accumulate_only=true: they add their instructions' estimates but must
1982+
// not trigger a threshold evaluation themselves. Deferring the check until
1983+
// control returns to the enclosing chain means a SCAN's collected outputs
1984+
// have been stacked into state and are reachable from the evaluation roots;
1985+
// evaluating inside the body could otherwise reset the shared counter while
1986+
// earlier outputs were still retained only in `collected`.
1987+
//
1988+
// accumulate_only suppresses only the threshold-triggered evaluation here.
1989+
// Op-internal evaluations are unaffected.
1990+
void run_chain(
1991+
const MLXProgram& prog,
1992+
uint32_t chain_idx,
1993+
ExecutionState& st,
1994+
StreamOrDevice stream,
1995+
size_t& pending_bytes,
1996+
bool accumulate_only) const {
19411997
if (chain_idx >= prog.instruction_chains.size()) {
19421998
throw std::runtime_error(
19431999
"run_chain: chain_idx " + std::to_string(chain_idx) +
19442000
" out of range (num_chains=" +
19452001
std::to_string(prog.instruction_chains.size()) + ")");
19462002
}
19472003
const auto& chain = prog.instruction_chains[chain_idx];
2004+
const size_t threshold = eval_threshold_bytes_;
19482005
size_t idx = 0;
19492006
for (const auto& instr : chain) {
19502007
st.begin_op(idx, op_name(instr.op));
19512008
if (instr.op == OpCode::SCAN) {
1952-
exec_scan(prog, std::get<ScanNode>(instr.node), st, stream);
2009+
exec_scan(
2010+
prog, std::get<ScanNode>(instr.node), st, stream, pending_bytes);
19532011
} else if (instr.op == OpCode::IF) {
1954-
exec_if(prog, std::get<IfNode>(instr.node), st, stream);
2012+
exec_if(prog, std::get<IfNode>(instr.node), st, stream, pending_bytes);
19552013
} else {
19562014
dispatch(instr, st, stream);
19572015
}
19582016
st.end_op();
19592017
++idx;
2018+
2019+
if (threshold != 0) {
2020+
// SCAN and IF already accumulated their own child instructions through
2021+
// the shared counter; charging the parent for them again would double
2022+
// count.
2023+
if (instr.op != OpCode::SCAN && instr.op != OpCode::IF) {
2024+
accumulate_instruction_bytes(instr, st, pending_bytes);
2025+
}
2026+
if (!accumulate_only && pending_bytes >= threshold) {
2027+
evaluate_state_tensors(st);
2028+
pending_bytes = 0;
2029+
}
2030+
}
19602031
}
19612032
}
19622033

19632034
private:
2035+
// Charge `pending_bytes` for one instruction. The estimate is the largest
2036+
// per-execution tensor the instruction touches, which tracks the size of what
2037+
// it just produced without needing to know which tid is its output.
2038+
// Constants and mutable buffers are excluded: they are not intermediates and
2039+
// evaluating does not release them.
2040+
//
2041+
// Only ever called when the mechanism is enabled.
2042+
static void accumulate_instruction_bytes(
2043+
const Instruction& instr,
2044+
const ExecutionState& st,
2045+
size_t& pending_bytes) {
2046+
accounting_calls_.fetch_add(1, std::memory_order_relaxed);
2047+
size_t widest = 0;
2048+
for_each_tid(instr, [&](Tid id) {
2049+
if (id.idx >= st.num_constants && !st.is_mutable_buffer(id)) {
2050+
uint32_t slot = st.tensor_index(id);
2051+
if (slot < st.tensors.size() && st.tensors[slot].has_value()) {
2052+
widest = std::max(widest, st.tensors[slot]->nbytes());
2053+
}
2054+
}
2055+
});
2056+
add_saturating(pending_bytes, widest);
2057+
}
2058+
2059+
// Saturating add: a pathological program must not wrap the counter back
2060+
// under the threshold and silently disable the mechanism.
2061+
static void add_saturating(size_t& acc, size_t add) {
2062+
if (add > std::numeric_limits<size_t>::max() - acc) {
2063+
acc = std::numeric_limits<size_t>::max();
2064+
} else {
2065+
acc += add;
2066+
}
2067+
}
2068+
2069+
// Materialize every live per-execution tensor, releasing the graph that
2070+
// produced them. Results are unchanged by evaluating early.
2071+
static void evaluate_state_tensors(ExecutionState& st) {
2072+
std::vector<::mlx::core::array> live;
2073+
live.reserve(st.tensors.size());
2074+
for (auto& t : st.tensors) {
2075+
if (t.has_value()) {
2076+
live.push_back(*t);
2077+
}
2078+
}
2079+
if (!live.empty()) {
2080+
::mlx::core::eval(live);
2081+
}
2082+
}
2083+
2084+
size_t eval_threshold_bytes_{0};
2085+
inline static std::atomic<uint64_t> accounting_calls_{0};
2086+
19642087
void exec_if(
19652088
const MLXProgram& prog,
19662089
const IfNode& n,
19672090
ExecutionState& st,
1968-
StreamOrDevice s) const {
2091+
StreamOrDevice s,
2092+
size_t& pending_bytes) const {
19692093
// Select one branch at runtime based on the integer condition.
19702094
// Nonzero -> then_chain, zero -> else_chain. The selected chain's
19712095
// instructions write the output slot(s) directly.
19722096
const int64_t cond = resolve_int(n.cond, st);
19732097
const uint32_t chain_idx =
19742098
(cond != 0) ? n.then_chain_idx : n.else_chain_idx;
1975-
run_chain(prog, chain_idx, st, s);
2099+
run_chain(prog, chain_idx, st, s, pending_bytes, /*accumulate_only=*/true);
19762100
}
19772101

19782102
void exec_scan(
19792103
const MLXProgram& prog,
19802104
const ScanNode& n,
19812105
ExecutionState& st,
1982-
StreamOrDevice s) const {
2106+
StreamOrDevice s,
2107+
size_t& pending_bytes) const {
19832108
int axis = n.scan_axis;
19842109
int T_int = st.const_tensor_ref(n.originals[0]).shape(axis);
19852110
size_t T = static_cast<size_t>(T_int);
@@ -2001,7 +2126,13 @@ class Interpreter {
20012126
s));
20022127
}
20032128

2004-
run_chain(prog, static_cast<uint32_t>(n.body_chain_idx), st, s);
2129+
run_chain(
2130+
prog,
2131+
static_cast<uint32_t>(n.body_chain_idx),
2132+
st,
2133+
s,
2134+
pending_bytes,
2135+
/*accumulate_only=*/true);
20052136

20062137
for (size_t i = 0; i < num_outputs; ++i) {
20072138
collected[i].push_back(st.const_tensor_ref(n.outputs[i]));
@@ -2011,6 +2142,21 @@ class Interpreter {
20112142
for (size_t i = 0; i < num_outputs; ++i) {
20122143
st.set_tensor(n.outputs[i], ::mlx::core::stack(collected[i], axis, s));
20132144
}
2145+
2146+
// The stacked outputs are new allocations the body's per-instruction
2147+
// estimates never saw, so charge for them here. Guarded like every other
2148+
// piece of the accounting.
2149+
if (eval_threshold_bytes_ != 0) {
2150+
for (size_t i = 0; i < num_outputs; ++i) {
2151+
const Tid id = n.outputs[i];
2152+
if (id.idx >= st.num_constants && !st.is_mutable_buffer(id)) {
2153+
uint32_t slot = st.tensor_index(id);
2154+
if (slot < st.tensors.size() && st.tensors[slot].has_value()) {
2155+
add_saturating(pending_bytes, st.tensors[slot]->nbytes());
2156+
}
2157+
}
2158+
}
2159+
}
20142160
}
20152161
void dispatch(const Instruction& instr, ExecutionState& st, StreamOrDevice s)
20162162
const {

‎backends/mlx/runtime/backend_options.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,36 @@ inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval";
4242
// errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle.
4343
inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init";
4444

45+
// Per-model runtime-spec key. Value N means: while running a method, evaluate
46+
// the live per-execution tensors once the intermediates produced since the last
47+
// evaluation exceed N bytes. 0/unset disables the mechanism entirely and is the
48+
// default.
49+
//
50+
// WHY: MLX is lazy. Interpreter::dispatch only builds graph nodes, and nothing
51+
// is materialized until MLXBackend::execute calls async_eval on the method
52+
// outputs, so for a long instruction chain every intermediate in the method is
53+
// live at the same instant. Whisper-small's 495-instruction encode peaks at
54+
// 1105 MB of MLX allocation against 95 MB of steady-state active memory, which
55+
// is what makes the model unusable on an iPhone (pytorch/executorch#22513).
56+
// Each evaluation costs a GPU sync, so the cost tracks the NUMBER of
57+
// evaluations; budgeting bytes rather than counting instructions puts them only
58+
// in the methods that actually allocate.
59+
//
60+
// NOTE that this is a THRESHOLD, not a hard memory limit. It is best-effort
61+
// evaluation scheduling, and peak footprint can exceed it:
62+
// - A long SCAN or IF branch accumulates across its whole body and is only
63+
// checked once control returns to the enclosing chain, so it can overshoot
64+
// by the size of that body.
65+
// - The per-instruction estimate is the largest tensor the instruction
66+
// touches, which can overcount (an op that only reads a large tensor is
67+
// charged for it) and so can trigger evaluation earlier than the true
68+
// pending bytes warrant.
69+
// - Ops that evaluate internally reduce the real pending work without
70+
// reducing the running estimate.
71+
// Treat it as a knob to trade GPU syncs against peak memory, and tune it
72+
// against measurements rather than expecting the value to bound RSS.
73+
inline constexpr char kEvalThresholdBytesKey[] = "eval_threshold_bytes";
74+
4575
} // namespace mlx
4676
} // namespace backends
4777
} // namespace executorch

‎backends/mlx/test/CMakeLists.txt‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,30 @@ if(EXECUTORCH_MLX_ENABLE_SANITIZERS)
115115
endif()
116116
add_test(NAME mlx_mutable_state COMMAND mlx_mutable_state_test)
117117

118+
# Lazy-graph evaluation threshold (kEvalThresholdBytesKey). Drives the
119+
# interpreter over hand-built programs; no model/tokenizer needed.
120+
et_cxx_test(
121+
mlx_eval_threshold_test
122+
SOURCES
123+
${CMAKE_CURRENT_LIST_DIR}/mlx_eval_threshold_test.cpp
124+
EXTRA_LIBS
125+
mlxdelegate
126+
mlx_schema
127+
mlx
128+
)
129+
target_include_directories(
130+
mlx_eval_threshold_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime
131+
)
132+
if(EXECUTORCH_MLX_ENABLE_SANITIZERS)
133+
target_compile_options(
134+
mlx_eval_threshold_test PRIVATE -fsanitize=address,undefined
135+
-fno-omit-frame-pointer
136+
)
137+
target_link_options(
138+
mlx_eval_threshold_test PRIVATE ${_mlx_sanitizer_link_options}
139+
)
140+
endif()
141+
118142
# Off-graph KV cache op-level test (no model/tokenizer needed). et_cxx_test
119143
# links GTest + executorch_core and registers the ctest target.
120144
et_cxx_test(

0 commit comments

Comments
 (0)