diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 4fdea010b40..0ce0f3b57a7 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -66,13 +66,17 @@ jobs: echo "::endgroup::" echo "::group::Build test runners" - ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test -j$(( $(sysctl -n hw.ncpu) - 1 )) + ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_flat_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) 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 flat KV-cache op test" + ./cmake-out/backends/mlx/test/mlx_flat_cache_test + echo "::endgroup::" + echo "::group::Run op unit tests" ${CONDA_RUN} python -m executorch.backends.mlx.test.run_all_tests -j4 --max-tasks-per-worker 10 --clean-after echo "::endgroup::" diff --git a/backends/mlx/runtime/MLXCacheImpl.h b/backends/mlx/runtime/MLXCacheImpl.h new file mode 100644 index 00000000000..b6f160e5c4a --- /dev/null +++ b/backends/mlx/runtime/MLXCacheImpl.h @@ -0,0 +1,52 @@ +/* + * 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. + */ + +#pragma once + +#include + +#include "MLXExecutor.h" // Tensor, StreamOrDevice + +namespace executorch { +namespace backends { +namespace mlx { + +// The K/V window to attend over + how to mask it; the cache owns the semantic. +// `kind` mirrors MLX SDPA's mask forms (no mask / "causal" / explicit tensor). +// Will be added: a `window` on Causal (sliding-window -- the +// mask_mod axis) and a `softcap`/bias field (ALiBi, softcap -- the score_mod +// axis). +struct AttendSpec { + Tensor K; + Tensor V; + enum class Mask { None, Causal, Explicit } kind; + std::optional mask; // Explicit only +}; + +// Op face of the off-graph KV cache; ExecutionState holds one (cross-cast from +// the registry's CacheBase*). Separate from CacheBase to avoid a diamond. +class MLXCacheImpl { + public: + virtual ~MLXCacheImpl() = default; + + // Write this step's K/V for `layer` at `position` (the run's logical start); + // return the window + mask kind. k/v are BHSD. `position` is a host int -- + // the caller reads it off the graph so the cache stays pure graph + integer + // bookkeeping. The cache owns the mask: a multi-token chain is Causal, a + // single decode token is None. + virtual AttendSpec update_and_fetch( + int layer, + int position, + const Tensor& k, + const Tensor& v, + StreamOrDevice s) = 0; +}; + +} // namespace mlx +} // namespace backends +} // namespace executorch diff --git a/backends/mlx/runtime/MLXExecutor.h b/backends/mlx/runtime/MLXExecutor.h index 978eaadabba..2c0bca4b27a 100644 --- a/backends/mlx/runtime/MLXExecutor.h +++ b/backends/mlx/runtime/MLXExecutor.h @@ -148,10 +148,14 @@ struct MutableBufferData { } }; +class MLXCacheImpl; // off-graph KV cache op face (MLXCacheImpl.h) + struct ExecutionState { const MLXProgram* program{nullptr}; const ConstantData* constants{nullptr}; // Shared, read-only MutableBufferData* mutable_buffers{nullptr}; // Per-handle, persistent + MLXCacheImpl* cache{ + nullptr}; // Per-session off-graph KV cache; survives reset() // Per-execution tensors: inputs, outputs, temps (NOT constants or mutable // buffers) diff --git a/backends/mlx/runtime/MLXFlatCache.h b/backends/mlx/runtime/MLXFlatCache.h new file mode 100644 index 00000000000..52bedc8f5d6 --- /dev/null +++ b/backends/mlx/runtime/MLXFlatCache.h @@ -0,0 +1,146 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include "MLXCacheImpl.h" // AttendSpec, MLXCacheImpl +#include "MLXExecutor.h" // to_shape + +#include +#include + +namespace executorch { +namespace backends { +namespace mlx { + +namespace cache = ::executorch::extension::llm::cache; + +// Per-layer K or V store, SDPA-major [1, H, capacity, D] (cells on axis 2), +// allocated at construction from config H/D/dtype. A write is a slice_update +// run (casting the update to the storage dtype if it differs); a read is a [0, +// len) slice. +class FlatPool { + public: + FlatPool(int capacity, int H, int D, ::mlx::core::Dtype dtype) + : dtype_(dtype), + buf_(::mlx::core::zeros( + to_shape(std::vector{1, H, capacity, D}), + dtype)) {} + + void write(int start, const Tensor& update, StreamOrDevice s) { + const int cap = static_cast(buf_.shape(2)); + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + const int T = static_cast(update.shape(2)); + if (start < 0 || start + T > cap) { + throw std::runtime_error("FlatPool::write: run out of bounds"); + } + if (static_cast(update.shape(1)) != H || + static_cast(update.shape(3)) != D) { + throw std::runtime_error("FlatPool::write: K/V heads/dim mismatch"); + } + const Tensor u = update.dtype() == dtype_ + ? update + : ::mlx::core::astype(update, dtype_, s); + std::vector lo{0, 0, start, 0}; + std::vector hi{1, H, start + T, D}; + buf_ = ::mlx::core::slice_update(buf_, u, to_shape(lo), to_shape(hi), s); + } + + Tensor read(int len, StreamOrDevice s) const { + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + std::vector lo{0, 0, 0, 0}; + std::vector hi{1, H, len, D}; + std::vector step{1, 1, 1, 1}; + return ::mlx::core::slice( + buf_, to_shape(lo), to_shape(hi), to_shape(step), s); + } + + private: + ::mlx::core::Dtype dtype_; + Tensor buf_; +}; + +// Single-sequence, full-history cache: the neutral SequenceCache bookkeeping +// over per-layer FlatPools. update_and_fetch plans the step (integer runs), +// writes the new K/V, reads the retained window, and declares the mask; the op +// handler owns q/scale and calls SDPA. +class MLXFlatCache : public cache::SequenceCache, public MLXCacheImpl { + public: + explicit MLXFlatCache(const cache::CacheConfig& cfg) + : cache::SequenceCache(cfg) { + if (!cfg.kv_dtype) { + throw std::runtime_error( + "MLXFlatCache: CacheConfig::kv_dtype (storage precision) must be set"); + } + const ::mlx::core::Dtype dt = + resolve_dtype(static_cast(*cfg.kv_dtype)); + kpool_.reserve(static_cast(cfg.n_layers)); + vpool_.reserve(static_cast(cfg.n_layers)); + for (int l = 0; l < cfg.n_layers; ++l) { + // layers size 1 = one config broadcast to every layer, else per-layer. + const cache::LayerConfig& lc = + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; + kpool_.emplace_back(cfg.capacity, lc.n_kv_heads, lc.head_dim, dt); + vpool_.emplace_back(cfg.capacity, lc.n_kv_heads, lc.head_dim, dt); + } + } + + AttendSpec update_and_fetch( + int layer, + int position, + const Tensor& k, + const Tensor& v, + StreamOrDevice s) override { + if (layer < 0 || layer >= static_cast(kpool_.size())) { + throw std::out_of_range("update_and_fetch: layer out of range"); + } + const int T = static_cast(k.shape(2)); // BHSD: seq axis is 2 + const int start = position; + + std::optional p = this->plan(layer, start, T); + if (!p) { + throw std::runtime_error( + "update_and_fetch: step exceeds capacity or invalid layer"); + } + // Flat is single-run: one write [start, start+T), one read [0, end). A + // wrapping (two-run) plan means a ring layer, which the ring cache handles; + // reject it here rather than silently drop the second run. + if (p->n_write != 1 || p->n_read != 1) { + throw std::runtime_error( + "update_and_fetch: expected a flat single-run plan"); + } + const size_t l = static_cast(layer); + kpool_[l].write(p->write[0].start, k, s); + vpool_[l].write(p->write[0].start, v, s); + Tensor K = kpool_[l].read(p->read[0].len, s); + Tensor V = vpool_[l].read(p->read[0].len, s); + this->commit(*p); + + // A multi-token chain is Causal (MLX "causal" is lower-right aligned, so + // fresh and chunked prefill are both correct with new tokens at the tail); + // a single decode token needs no mask. + const AttendSpec::Mask kind = + (T > 1) ? AttendSpec::Mask::Causal : AttendSpec::Mask::None; + return AttendSpec{K, V, kind, std::nullopt}; + } + + private: + std::vector kpool_; + std::vector vpool_; +}; + +} // namespace mlx +} // namespace backends +} // namespace executorch diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 3c3c2c323a8..5de29233a22 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -8,6 +8,7 @@ #pragma once +#include "MLXCacheImpl.h" #include "MLXExecutor.h" #include @@ -293,6 +294,50 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor(n.out, std::move(out)); } +inline void exec_update_and_attend( + const UpdateAndAttendNode& n, + ExecutionState& st, + StreamOrDevice s) { + if (!st.cache) { + throw std::runtime_error("update_and_attend: no cache installed"); + } + // The cache does the KV write + read and declares the mask; the handler owns + // the query side (q, scale) and calls SDPA. + const array& q = st.const_tensor_ref(n.q); + // The run's start is position[0]. Read it host-side here -- a tiny + // device->host sync -- so the cache stays pure graph + integer bookkeeping. + // If position later arrives as a host-side constant, this is the single spot + // to change. + array pos = astype(st.const_tensor_ref(n.position), int32, s); + eval(pos); + const int position = pos.data()[0]; + AttendSpec spec = st.cache->update_and_fetch( + n.layer_id, + position, + st.const_tensor_ref(n.k), + st.const_tensor_ref(n.v), + s); + // Match stored K/V to the query dtype before SDPA (no-op when equal; the + // storage precision may differ from the compute dtype). + array K = spec.K.dtype() == q.dtype() ? spec.K : astype(spec.K, q.dtype(), s); + array V = spec.V.dtype() == q.dtype() ? spec.V : astype(spec.V, q.dtype(), s); + std::string mask_mode = spec.kind == AttendSpec::Mask::Causal ? "causal" : ""; + array out = fast::scaled_dot_product_attention( + q, + K, + V, + static_cast(n.scale), + mask_mode, + spec.mask, + std::nullopt, + s); + // Honor the op's output-dtype contract (unset -> SDPA's native output). + if (n.out_dtype) { + out = astype(out, resolve_dtype(*n.out_dtype), s); + } + st.set_tensor(n.out, std::move(out)); +} + inline void exec_add(const AddNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor( n.out, add(st.const_tensor_ref(n.a), st.const_tensor_ref(n.b), s)); @@ -1959,6 +2004,10 @@ class Interpreter { case OpCode::SDPA: ops::exec_sdpa(std::get(instr.node), st, s); break; + case OpCode::UPDATE_AND_ATTEND: + ops::exec_update_and_attend( + std::get(instr.node), st, s); + break; case OpCode::ADD: ops::exec_add(std::get(instr.node), st, s); break; diff --git a/backends/mlx/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs index 281199a8002..0f4dfe0687a 100644 --- a/backends/mlx/serialization/schema.fbs +++ b/backends/mlx/serialization/schema.fbs @@ -153,6 +153,17 @@ table SdpaNode { causal: bool = false; } +table UpdateAndAttendNode { + q: Tid (required); + k: Tid (required); + v: Tid (required); + position: Tid (required); + out: Tid (required); + layer_id: int32; + scale: float; + out_dtype: int8 = null; // ET ScalarType; null = keep SDPA's native output +} + table AddNode { a: Tid (required); b: Tid (required); @@ -1170,7 +1181,8 @@ union OpNode { BitwiseOrNode, BitwiseXorNode, IfNode, - RandomBitsNode + RandomBitsNode, + UpdateAndAttendNode // BC: Add new op nodes here (append only) } diff --git a/backends/mlx/test/CMakeLists.txt b/backends/mlx/test/CMakeLists.txt index 2d494652138..6fa0ad50b39 100644 --- a/backends/mlx/test/CMakeLists.txt +++ b/backends/mlx/test/CMakeLists.txt @@ -88,3 +88,27 @@ if(EXECUTORCH_MLX_ENABLE_SANITIZERS) ) endif() add_test(NAME mlx_mutable_state COMMAND mlx_mutable_state_test) + +# Off-graph flat KV cache op-level test (no model/tokenizer needed). et_cxx_test +# links GTest + executorch_core and registers the ctest target. +et_cxx_test( + mlx_flat_cache_test + SOURCES + ${CMAKE_CURRENT_LIST_DIR}/mlx_flat_cache_test.cpp + EXTRA_LIBS + mlxdelegate + mlx_schema + mlx +) +target_include_directories( + mlx_flat_cache_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime +) +if(EXECUTORCH_MLX_ENABLE_SANITIZERS) + target_compile_options( + mlx_flat_cache_test PRIVATE -fsanitize=address,undefined + -fno-omit-frame-pointer + ) + target_link_options( + mlx_flat_cache_test PRIVATE ${_mlx_sanitizer_link_options} + ) +endif() diff --git a/backends/mlx/test/mlx_flat_cache_test.cpp b/backends/mlx/test/mlx_flat_cache_test.cpp new file mode 100644 index 00000000000..bc9ccb03ea6 --- /dev/null +++ b/backends/mlx/test/mlx_flat_cache_test.cpp @@ -0,0 +1,184 @@ +/* + * 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. + */ + +// Op-level test for the off-graph flat KV cache (MLXFlatCache / FlatPool). +// +// Drives MLXFlatCache::update_and_fetch directly (no interpreter / .pte), then +// attends the returned AttendSpec exactly as the op handler does, and checks +// the result against MLX SDPA over the full K/V history the cache should have +// assembled -- verifying the plan/write/read window and mask kind across +// prefill (Causal) and decode (None), plus the capacity-reject, storage-dtype, +// and required-dtype paths. +// +// Must run on Apple Silicon: MLX SDPA needs the Metal backend. + +#include "MLXFlatCache.h" + +#include + +#include + +#include +#include +#include +#include + +using namespace ::executorch::backends::mlx; +namespace cache = ::executorch::extension::llm::cache; +using ::mlx::core::array; + +namespace { + +// Max absolute difference within tolerance. Computed in float32: item() +// reads sizeof(float) bytes, so calling it on an fp16 scalar misreads the +// buffer. +bool allclose(const array& a, const array& b, float atol) { + using namespace ::mlx::core; + array m = max(abs(subtract(astype(a, float32), astype(b, float32)))); + eval(m); + return m.item() <= atol; +} + +cache::CacheConfig flat_config( + int capacity, + int n_layers, + int n_kv_heads, + int head_dim, + std::optional kv_dtype = std::nullopt) { + cache::CacheConfig cfg; + cfg.capacity = capacity; + cfg.n_layers = n_layers; + cfg.layers = {cache::LayerConfig{ + cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}, + n_kv_heads, + head_dim}}; + cfg.kv_dtype = kv_dtype; + return cfg; +} + +class MLXFlatCacheTest : public ::testing::Test { + protected: + const int H = 2; + const int D = 8; + const float scale = 1.0f / std::sqrt(static_cast(D)); + ::mlx::core::StreamOrDevice s = {}; + + array randn(int T, ::mlx::core::Dtype dt) { + return ::mlx::core::random::normal( + to_shape(std::vector{1, H, T, D}), dt); + } +}; + +// Prefill: T=4 at position 0 -> Causal (lower-right aligned). +TEST_F(MLXFlatCacheTest, PrefillIsCausal) { + using namespace ::mlx::core; + MLXFlatCache c(flat_config( + /*capacity=*/32, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + const int T0 = 4; + array q0 = randn(T0, float16); + array k0 = randn(T0, float16); + array v0 = randn(T0, float16); + + AttendSpec spec0 = c.update_and_fetch(0, /*position=*/0, k0, v0, s); + EXPECT_EQ(spec0.kind, AttendSpec::Mask::Causal); + + array cand0 = fast::scaled_dot_product_attention( + q0, + spec0.K, + spec0.V, + scale, + std::string("causal"), + std::nullopt, + std::nullopt, + s); + array ref0 = fast::scaled_dot_product_attention( + q0, k0, v0, scale, std::string("causal"), std::nullopt, std::nullopt, s); + EXPECT_TRUE(allclose(cand0, ref0, 1e-2f)); +} + +// Decode: after a T=4 prefill, a single query at position 4 -> None and must +// attend the full assembled history. +TEST_F(MLXFlatCacheTest, DecodeAttendsFullHistory) { + using namespace ::mlx::core; + MLXFlatCache c(flat_config( + /*capacity=*/32, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + const int T0 = 4; + array k0 = randn(T0, float16); + array v0 = randn(T0, float16); + c.update_and_fetch(0, /*position=*/0, k0, v0, s); // prefill + + array q1 = randn(1, float16); + array k1 = randn(1, float16); + array v1 = randn(1, float16); + AttendSpec spec1 = c.update_and_fetch(0, /*position=*/T0, k1, v1, s); + EXPECT_EQ(spec1.kind, AttendSpec::Mask::None); + + array cand1 = fast::scaled_dot_product_attention( + q1, + spec1.K, + spec1.V, + scale, + std::string(""), + std::nullopt, + std::nullopt, + s); + array Khist = concatenate(std::vector{k0, k1}, 2, s); + array Vhist = concatenate(std::vector{v0, v1}, 2, s); + array ref1 = fast::scaled_dot_product_attention( + q1, Khist, Vhist, scale, std::string(""), std::nullopt, std::nullopt, s); + EXPECT_TRUE(allclose(cand1, ref1, 1e-2f)); +} + +// A step past capacity is rejected (plan returns nullopt). +TEST_F(MLXFlatCacheTest, StepPastCapacityThrows) { + using namespace ::mlx::core; + MLXFlatCache c(flat_config( + /*capacity=*/32, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + array kx = randn(1, float16); + EXPECT_ANY_THROW(c.update_and_fetch(0, /*position=*/32, kx, kx, s)); +} + +// Storage dtype != compute: fp32 input, fp16 storage. The cache casts on write, +// so the read-back K/V are exactly the fp16 of the input. +TEST_F(MLXFlatCacheTest, StorageDtypeDiffersCastsOnWrite) { + using namespace ::mlx::core; + MLXFlatCache c16(flat_config( + /*capacity=*/32, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + const int T0 = 4; + array k2 = randn(T0, float32); + array v2 = randn(T0, float32); + AttendSpec spec2 = c16.update_and_fetch(0, /*position=*/0, k2, v2, s); + EXPECT_EQ(spec2.K.dtype(), float16); + EXPECT_EQ(spec2.V.dtype(), float16); + EXPECT_TRUE(allclose(spec2.K, astype(k2, float16, s), 0.0f)); + EXPECT_TRUE(allclose(spec2.V, astype(v2, float16, s), 0.0f)); +} + +// The MLX cache requires an explicit storage dtype. +TEST_F(MLXFlatCacheTest, UnsetKvDtypeThrows) { + cache::CacheConfig cfg = flat_config(/*capacity=*/32, /*n_layers=*/1, H, D); + EXPECT_ANY_THROW(MLXFlatCache{cfg}); +} + +} // namespace diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 3ab424830e0..c657571a490 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -127,6 +127,9 @@ struct CacheConfig { std::vector layers; int initial_capacity = 512; std::optional max_write; + // ET ScalarType storage precision; a backend that stores K/V may require it + // (MLX raises if unset). + std::optional kv_dtype; }; } // namespace cache