diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index ec2d5b41818..f9a9a1f0dd2 100644 --- a/.github/workflows/mlx.yml +++ b/.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 diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index 0689388412f..9acb9228098 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/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(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(bytes)); + } + if (!processed || !processed->data() || processed->size() == 0) { throw std::runtime_error("init: null or empty delegate payload"); } diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 498c0e34f7c..f13d827d4ab 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -19,6 +19,11 @@ #include #include +#include +#include +#include +#include + namespace executorch { namespace backends { namespace mlx { @@ -1926,6 +1931,29 @@ 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, @@ -1933,11 +1961,39 @@ class Interpreter { 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) + @@ -1945,41 +2001,110 @@ class Interpreter { 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(instr.node), st, stream); + exec_scan( + prog, std::get(instr.node), st, stream, pending_bytes); } else if (instr.op == OpCode::IF) { - exec_if(prog, std::get(instr.node), st, stream); + exec_if(prog, std::get(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::max() - acc) { + acc = std::numeric_limits::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 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(T_int); @@ -2001,7 +2126,13 @@ class Interpreter { s)); } - run_chain(prog, static_cast(n.body_chain_idx), st, s); + run_chain( + prog, + static_cast(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 { diff --git a/backends/mlx/runtime/backend_options.h b/backends/mlx/runtime/backend_options.h index 2f96ce93525..1217c58e0d9 100644 --- a/backends/mlx/runtime/backend_options.h +++ b/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 diff --git a/backends/mlx/test/CMakeLists.txt b/backends/mlx/test/CMakeLists.txt index 36c7cc463bf..310f1ca5fc6 100644 --- a/backends/mlx/test/CMakeLists.txt +++ b/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( diff --git a/backends/mlx/test/mlx_eval_threshold_test.cpp b/backends/mlx/test/mlx_eval_threshold_test.cpp new file mode 100644 index 00000000000..41670a96145 --- /dev/null +++ b/backends/mlx/test/mlx_eval_threshold_test.cpp @@ -0,0 +1,275 @@ +/* + * 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. + */ + +// Tests for the lazy-graph evaluation threshold (kEvalThresholdBytesKey). +// +// Drives Interpreter::run_chain directly over hand-built programs (no .pte), +// covering: that the disabled path does no accounting at all, that nested +// IF/SCAN chains accumulate through the caller's counter without being +// double-counted, that crossing the threshold actually materializes the live +// tensors mid-chain, and that results are unchanged by any of it. +// +// Must run on Apple Silicon: MLX needs the Metal backend. + +#include "MLXInterpreter.h" + +#include + +#include + +#include + +using namespace ::executorch::backends::mlx; +using ::mlx::core::array; + +namespace { + +constexpr uint32_t kIn = 0; // input tid +constexpr uint32_t kOut = 1; // output tid +constexpr uint32_t kTemp0 = 2; // first temp tid + +Instruction make_add(uint32_t a, uint32_t b, uint32_t out) { + Instruction instr; + instr.op = OpCode::ADD; + AddNode node; + node.a = Tid{a}; + node.b = Tid{b}; + node.out = Tid{out}; + instr.node = node; + return instr; +} + +// A chain of `n` ADDs: in -> temp0 -> temp1 -> ... -> out. +std::vector make_add_chain(uint32_t n) { + std::vector chain; + uint32_t prev = kIn; + for (uint32_t i = 0; i < n; ++i) { + const bool last = (i + 1 == n); + const uint32_t out = last ? kOut : (kTemp0 + i); + chain.push_back(make_add(prev, prev, out)); + prev = out; + } + return chain; +} + +// Program whose main chain is `n` ADDs and nothing else. +MLXProgram make_flat_program(uint32_t n) { + MLXProgram program; + program.num_input_tensors = 1; + program.num_output_tensors = 1; + program.num_temp_tensors = n; // generous; unused slots stay nullopt + program.instruction_chains.push_back(make_add_chain(n)); + program.main_chain_idx = 0; + return program; +} + +// Program whose main chain is a single IF; the taken branch is `n` ADDs. +// Exercises nested accumulation: the branch's instructions must be charged to +// the caller's counter, and the IF itself must not be charged again. +MLXProgram make_if_program(uint32_t n) { + MLXProgram program; + program.num_input_tensors = 1; + program.num_output_tensors = 1; + program.num_temp_tensors = n; + + Instruction if_instr; + if_instr.op = OpCode::IF; + IfNode node; + node.cond = static_cast(1); // always take then_chain + node.then_chain_idx = 1; + node.else_chain_idx = 2; + if_instr.node = node; + + program.instruction_chains.push_back({if_instr}); // chain 0: main + program.instruction_chains.push_back(make_add_chain(n)); // chain 1: then + program.instruction_chains.push_back(make_add_chain(1)); // chain 2: else + program.main_chain_idx = 0; + return program; +} + +// Bind a state for `program` and seed the input with a non-trivial tensor. +// `mb` sizes the input so byte thresholds are easy to reason about. +void bind_state( + ExecutionState& st, + const MLXProgram& program, + ConstantData& constants, + MutableBufferData& bufs, + int floats) { + st.bind(program, constants, bufs); + st.set_tensor( + Tid{kIn}, ::mlx::core::full({floats}, 1.0f, ::mlx::core::float32)); +} + +// Bytes of one input tensor of `floats` elements. +size_t input_bytes(int floats) { + return static_cast(floats) * sizeof(float); +} + +std::vector read_output(ExecutionState& st) { + const array& out = st.tensors[st.tensor_index(Tid{kOut})].value(); + ::mlx::core::eval(out); + return std::vector(out.data(), out.data() + out.size()); +} + +} // namespace + +// The whole mechanism must be inert when the option is unset. Not merely "no +// evaluation happens": no traversal, no nbytes() queries, no accumulation. +TEST(MLXEvalThreshold, DisabledDoesNoAccountingAtAll) { + const uint32_t kN = 8; + MLXProgram program = make_flat_program(kN); + ConstantData constants; + MutableBufferData bufs; + ExecutionState st; + bind_state(st, program, constants, bufs, 1024); + + Interpreter interp; // default: disabled + ASSERT_EQ(interp.eval_threshold_bytes(), 0u); + + Interpreter::reset_accounting_calls(); + interp.run(program, st); + EXPECT_EQ(Interpreter::accounting_calls(), 0u); +} + +// Enabled but with a threshold no run can reach: every instruction is still +// accounted, and nothing is evaluated early. +TEST(MLXEvalThreshold, EnabledAccountsEveryInstruction) { + const uint32_t kN = 8; + MLXProgram program = make_flat_program(kN); + ConstantData constants; + MutableBufferData bufs; + ExecutionState st; + bind_state(st, program, constants, bufs, 1024); + + Interpreter interp; + interp.set_eval_threshold_bytes(std::numeric_limits::max()); + EXPECT_EQ(interp.eval_threshold_bytes(), std::numeric_limits::max()); + + Interpreter::reset_accounting_calls(); + interp.run(program, st); + EXPECT_EQ(Interpreter::accounting_calls(), kN); + for (uint32_t i = 0; i < kN - 1; ++i) { + const array& temp = st.tensors[st.tensor_index(Tid{kTemp0 + i})].value(); + EXPECT_FALSE(temp.is_available()) << "temp=" << i; + } + const array& out = st.tensors[st.tensor_index(Tid{kOut})].value(); + EXPECT_FALSE(out.is_available()); +} + +// An IF's branch instructions are charged to the caller's counter, and the IF +// instruction itself is not charged on top of them. +TEST(MLXEvalThreshold, NestedIfAccumulatesWithoutDoubleCounting) { + const uint32_t kN = 6; + MLXProgram program = make_if_program(kN); + ConstantData constants; + MutableBufferData bufs; + ExecutionState st; + bind_state(st, program, constants, bufs, 1024); + + Interpreter interp; + interp.set_eval_threshold_bytes(std::numeric_limits::max()); + + Interpreter::reset_accounting_calls(); + interp.run(program, st); + // kN branch instructions, and NOT kN + 1: the IF is excluded at the parent. + EXPECT_EQ(Interpreter::accounting_calls(), kN); +} + +// Crossing the threshold must actually materialize the live tensors. Without a +// forced evaluation an MLX array built by dispatch alone is not available. +TEST(MLXEvalThreshold, CrossingThresholdEvaluatesLiveTensors) { + const int kFloats = 4096; + const uint32_t kN = 9; + MLXProgram program = make_flat_program(kN); + ConstantData constants; + MutableBufferData bufs; + + // Threshold of two instructions' worth: evaluation must fire mid-chain. + { + ExecutionState st; + bind_state(st, program, constants, bufs, kFloats); + Interpreter interp; + interp.set_eval_threshold_bytes(2 * input_bytes(kFloats)); + interp.run(program, st); + // The eighth ADD triggers a barrier; the ninth stays below the threshold. + for (uint32_t i = 0; i < kN - 1; ++i) { + const array& temp = st.tensors[st.tensor_index(Tid{kTemp0 + i})].value(); + EXPECT_TRUE(temp.is_available()) << "temp=" << i; + } + const array& out = st.tensors[st.tensor_index(Tid{kOut})].value(); + EXPECT_FALSE(out.is_available()); + } + + // Disabled: the same slot is still an unevaluated graph node. + { + ExecutionState st; + bind_state(st, program, constants, bufs, kFloats); + Interpreter interp; + interp.run(program, st); + const array& first_temp = st.tensors[st.tensor_index(Tid{kTemp0})].value(); + EXPECT_FALSE(first_temp.is_available()); + } +} + +// Evaluating early must not change results, on a flat chain or through a +// nested one. +TEST(MLXEvalThreshold, OutputsAreUnchangedByTheThreshold) { + const int kFloats = 2048; + const uint32_t kN = 8; + + for (bool nested : {false, true}) { + MLXProgram program = nested ? make_if_program(kN) : make_flat_program(kN); + ConstantData constants; + MutableBufferData bufs; + + ExecutionState off; + bind_state(off, program, constants, bufs, kFloats); + Interpreter disabled; + disabled.run(program, off); + const std::vector expected = read_output(off); + + ExecutionState on; + bind_state(on, program, constants, bufs, kFloats); + Interpreter enabled; + enabled.set_eval_threshold_bytes(input_bytes(kFloats)); // fires often + enabled.run(program, on); + const std::vector actual = read_output(on); + + ASSERT_EQ(actual.size(), expected.size()) << "nested=" << nested; + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(actual[i], expected[i]) << "nested=" << nested << " i=" << i; + } + } +} + +// The setting lives on the interpreter, so two handles configured differently +// do not interfere. +TEST(MLXEvalThreshold, SettingsArePerInterpreter) { + MLXProgram program = make_flat_program(4); + ConstantData constants; + MutableBufferData bufs; + + Interpreter enabled; + enabled.set_eval_threshold_bytes(1024); + Interpreter disabled; + + EXPECT_EQ(enabled.eval_threshold_bytes(), 1024u); + EXPECT_EQ(disabled.eval_threshold_bytes(), 0u); + + ExecutionState st_off; + bind_state(st_off, program, constants, bufs, 256); + Interpreter::reset_accounting_calls(); + disabled.run(program, st_off); + EXPECT_EQ(Interpreter::accounting_calls(), 0u); + + ExecutionState st_on; + bind_state(st_on, program, constants, bufs, 256); + Interpreter::reset_accounting_calls(); + enabled.run(program, st_on); + EXPECT_GT(Interpreter::accounting_calls(), 0u); +} diff --git a/docs/source/backends/mlx/mlx-overview.md b/docs/source/backends/mlx/mlx-overview.md index 643cd1d4190..f9b9a8dc722 100644 --- a/docs/source/backends/mlx/mlx-overview.md +++ b/docs/source/backends/mlx/mlx-overview.md @@ -141,6 +141,61 @@ There is also an `mlx-debug` preset useful during development: cmake --workflow --preset mlx-debug ``` +## Runtime Options + +The MLX backend reads optional per-model runtime specs, set through a +`LoadBackendOptionsMap` keyed by the backend id `MLXBackend`. All are optional +and off by default. + +### `eval_threshold_bytes` (int) + +MLX is lazy: dispatching an instruction only builds a graph node, and nothing is +materialized until the method's outputs are evaluated. For a long instruction +chain that means every intermediate in the method is live at the same instant, +so peak memory tracks the size of the whole graph rather than the working set. +Whisper-small's 495-instruction `encode` peaks at 1105 MB of MLX allocation +against 95 MB of steady-state active memory. + +Set this key to N to evaluate the live per-execution tensors once the +intermediates produced since the last evaluation exceed N bytes. Each evaluation +costs a GPU sync, so the cost tracks the *number* of evaluations, and budgeting +bytes rather than instructions puts them only in the methods that actually +allocate. + +`0` (the default) disables the mechanism entirely, preserving the previous +behaviour with no accounting overhead of any kind. + +```cpp +#include +#include + +executorch::runtime::BackendOptions<1> opts; +opts.set_option(executorch::backends::mlx::kEvalThresholdBytesKey, + 512 * 1024 * 1024); +``` + +Measured on an iPhone 16, whisper-small int8, full pipeline, medians of +interleaved rounds: + +| setting | peak MB | peak while loaded | pipeline ms | +| --- | --- | --- | --- | +| `0` (disabled) | 1194.4 | 763.7 | 885.2 | +| `512 MB` | 692.8 | 261.0 | 831.3 | + +This is a **threshold, not a hard memory limit**. It is best-effort evaluation +scheduling and peak footprint can exceed the value: + +- 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 evaluate earlier than the true pending bytes warrant. +- Ops that evaluate internally reduce the real pending work without reducing + the running estimate. + +Tune it against measurements rather than expecting the value to bound RSS. + ## Reference **→{doc}`/backends/mlx/mlx-troubleshooting` — Debug common issues.**