diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 4fdea010b40..800bc91419f 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_sequence_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 KV-cache op test" + ./cmake-out/backends/mlx/test/mlx_sequence_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/MLXCache.h b/backends/mlx/runtime/MLXCache.h new file mode 100644 index 00000000000..8c4e45d8cbe --- /dev/null +++ b/backends/mlx/runtime/MLXCache.h @@ -0,0 +1,55 @@ +/* + * 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 +}; + +// Tensor-typed op face of the off-graph KV cache, kept separate from the +// neutral CacheBase (which is tensor-free) so a cache can expose both without a +// diamond. ExecutionState holds one; nothing assigns it yet -- the registry +// that owns the cache and hands this pointer to the executor lands in a +// follow-up, until which exec_update_and_attend is unreachable. +class MLXCache { + public: + virtual ~MLXCache() = 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..eea49e40eb5 100644 --- a/backends/mlx/runtime/MLXExecutor.h +++ b/backends/mlx/runtime/MLXExecutor.h @@ -148,10 +148,13 @@ struct MutableBufferData { } }; +class MLXCache; // off-graph KV cache op face (MLXCache.h) + struct ExecutionState { const MLXProgram* program{nullptr}; const ConstantData* constants{nullptr}; // Shared, read-only MutableBufferData* mutable_buffers{nullptr}; // Per-handle, persistent + MLXCache* 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/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 3c3c2c323a8..e85ae894d80 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -8,6 +8,7 @@ #pragma once +#include "MLXCache.h" #include "MLXExecutor.h" #include @@ -293,6 +294,71 @@ 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"); + } + if (!n.scale) { + throw std::runtime_error("update_and_attend: scale is not set"); + } + // 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 host-side so the cache stays pure + // graph + integer bookkeeping. This is a device->host sync per node, i.e. + // n_layers stalls per step -- not one -- even though every layer of a step + // sees the same position. The fix is for position to arrive as a host-side + // constant instead of a graph tensor; 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); + // MLX takes the mask as a mode string plus an optional tensor. Switch rather + // than test for Causal: None and Explicit both map to "" and are told apart + // only by spec.mask, so an Explicit with no mask would silently attend + // unmasked. + std::string mask_mode; + switch (spec.kind) { + case AttendSpec::Mask::None: + break; + case AttendSpec::Mask::Causal: + mask_mode = "causal"; + break; + case AttendSpec::Mask::Explicit: + if (!spec.mask) { + throw std::runtime_error( + "update_and_attend: Explicit mask kind with no mask tensor"); + } + break; + } + 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 +2025,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/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h new file mode 100644 index 00000000000..1c1fabe8054 --- /dev/null +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -0,0 +1,170 @@ +/* + * 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 "MLXCache.h" // AttendSpec, MLXCache +#include "MLXExecutor.h" // resolve_dtype + +#include +#include + +namespace executorch { +namespace backends { +namespace mlx { + +namespace cache = ::executorch::extension::llm::cache; + +// Per-layer K or V store, SDPA-major [1, H, slots, D] (cells on axis 2). The +// planner hands down physical runs (it has already applied any ring modulo), so +// the pool is layout-agnostic: flat and ring differ only in how many slots the +// layer asks for and how many runs a step produces. +class Pool { + public: + Pool(int slots, int H, int D, ::mlx::core::Dtype dtype) + : dtype_(dtype), + buf_(::mlx::core::zeros(::mlx::core::Shape{1, H, slots, D}, dtype)) {} + + // Place `update` at the run's physical start, casting to the storage dtype if + // it differs. + // + // Each write chains a lazy slice_update onto the previous buf_. That stays + // bounded only because the caller eval()s the step's output, which + // transitively materializes buf_; a caller that never eval()s would grow the + // graph without limit. + void write(const cache::Run& run, const Tensor& update, StreamOrDevice s) { + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + if (run.start < 0 || run.start + run.len > slots()) { + throw std::runtime_error("Pool::write: run out of bounds"); + } + if (static_cast(update.shape(2)) != run.len) { + throw std::runtime_error("Pool::write: update length != run length"); + } + if (static_cast(update.shape(1)) != H || + static_cast(update.shape(3)) != D) { + throw std::runtime_error("Pool::write: K/V heads/dim mismatch"); + } + const Tensor u = update.dtype() == dtype_ + ? update + : ::mlx::core::astype(update, dtype_, s); + buf_ = ::mlx::core::slice_update( + buf_, + u, + ::mlx::core::Shape{0, 0, run.start, 0}, + ::mlx::core::Shape{1, H, run.start + run.len, D}, + s); + } + + // The run's cells, [start, start+len). Ring reads start mid-pool, so the run + // start matters here as much as it does for a write. + Tensor read(const cache::Run& run, StreamOrDevice s) const { + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + if (run.start < 0 || run.start + run.len > slots()) { + throw std::runtime_error("Pool::read: run out of bounds"); + } + return ::mlx::core::slice( + buf_, + ::mlx::core::Shape{0, 0, run.start, 0}, + ::mlx::core::Shape{1, H, run.start + run.len, D}, + ::mlx::core::Shape{1, 1, 1, 1}, + s); + } + + int slots() const { + return static_cast(buf_.shape(2)); + } + + private: + ::mlx::core::Dtype dtype_; + Tensor buf_; +}; + +// The MLX byte layer behind the neutral SequenceCache: one Pool per layer for K +// and one for V, sized from that layer's own policy. 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. +// +// Only flat layers are wired up so far. Ring layers are rejected at +// construction rather than silently mis-served: the plumbing below handles +// their runs, but the windowed mask and read_base_pos (RoPE alignment) are not +// implemented yet. +class MLXSequenceCache : public cache::SequenceCache, public MLXCache { + public: + explicit MLXSequenceCache(const cache::CacheConfig& cfg) + : cache::SequenceCache(cfg) { + 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]; + if (lc.policy.kind != cache::LayerPolicy::Kind::Flat) { + throw std::runtime_error( + "MLXSequenceCache: only Flat layers are supported so far"); + } + // Flat retains all history, so the layer's pool is the full cap. A ring + // layer will instead ask for window + max_write - 1 slots. + const int slots = cfg.capacity; + kpool_.emplace_back(slots, lc.n_kv_heads, lc.head_dim, dt); + vpool_.emplace_back(slots, 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 + + std::optional p = this->plan(layer, position, T); + if (!p) { + throw std::runtime_error( + "update_and_fetch: step exceeds capacity or invalid layer"); + } + // Flat layers never wrap, and ring layers are refused at construction, so + // every accepted step is a single run. Ring's split/rejoin lands with ring. + if (p->n_write != 1 || p->n_read != 1) { + throw std::runtime_error("update_and_fetch: expected a single-run plan"); + } + const size_t l = static_cast(layer); + kpool_[l].write(p->write[0], k, s); + vpool_[l].write(p->write[0], v, s); + Tensor K = kpool_[l].read(p->read[0], s); + Tensor V = vpool_[l].read(p->read[0], 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/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs index 281199a8002..99be50292db 100644 --- a/backends/mlx/serialization/schema.fbs +++ b/backends/mlx/serialization/schema.fbs @@ -153,6 +153,18 @@ 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 = null; // required in practice; null so an omitted scale + // fails loudly instead of defaulting to 0 + out_dtype: int8 = null; // ET ScalarType; null = keep SDPA's native output +} + table AddNode { a: Tid (required); b: Tid (required); @@ -1170,7 +1182,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..9b9d5cb4f85 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 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_sequence_cache_test + SOURCES + ${CMAKE_CURRENT_LIST_DIR}/mlx_sequence_cache_test.cpp + EXTRA_LIBS + mlxdelegate + mlx_schema + mlx +) +target_include_directories( + mlx_sequence_cache_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime +) +if(EXECUTORCH_MLX_ENABLE_SANITIZERS) + target_compile_options( + mlx_sequence_cache_test PRIVATE -fsanitize=address,undefined + -fno-omit-frame-pointer + ) + target_link_options( + mlx_sequence_cache_test PRIVATE ${_mlx_sanitizer_link_options} + ) +endif() diff --git a/backends/mlx/test/mlx_sequence_cache_test.cpp b/backends/mlx/test/mlx_sequence_cache_test.cpp new file mode 100644 index 00000000000..34a5fa40d7a --- /dev/null +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -0,0 +1,165 @@ +/* + * 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 KV cache (MLXSequenceCache / Pool). +// +// Drives MLXSequenceCache::update_and_fetch directly (no interpreter / .pte) +// and checks the AttendSpec it returns against the K/V history the cache should +// have assembled -- verifying the plan/write/read window and the mask kind +// across prefill (Causal) and decode (None), plus the capacity-reject and +// storage-dtype paths. The window is compared directly rather than through +// SDPA: attending both sides would only test equality through a lossy kernel. +// +// Must run on Apple Silicon: MLX needs the Metal backend. + +#include "MLXSequenceCache.h" + +#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, + int kv_dtype) { + 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 MLXSequenceCacheTest : public ::testing::Test { + protected: + const int H = 2; + const int D = 8; + ::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(MLXSequenceCacheTest, PrefillIsCausal) { + using namespace ::mlx::core; + MLXSequenceCache 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); + + AttendSpec spec0 = c.update_and_fetch(0, /*position=*/0, k0, v0, s); + EXPECT_EQ(spec0.kind, AttendSpec::Mask::Causal); + EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); + EXPECT_TRUE(allclose(spec0.V, v0, 0.0f)); +} + +// Decode: after a T=4 prefill, a single token at position 4 -> None, and the +// window is the full assembled history (prefill ++ the new token). +TEST_F(MLXSequenceCacheTest, DecodeReadsFullHistory) { + using namespace ::mlx::core; + MLXSequenceCache 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 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); + EXPECT_TRUE( + allclose(spec1.K, concatenate(std::vector{k0, k1}, 2, s), 0.0f)); + EXPECT_TRUE( + allclose(spec1.V, concatenate(std::vector{v0, v1}, 2, s), 0.0f)); +} + +// A step past capacity is rejected (plan returns nullopt). +TEST_F(MLXSequenceCacheTest, StepPastCapacityThrows) { + using namespace ::mlx::core; + MLXSequenceCache 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(MLXSequenceCacheTest, StorageDtypeDiffersCastsOnWrite) { + using namespace ::mlx::core; + MLXSequenceCache 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)); +} + +// A run is placed and fetched at its own physical start. Flat runs always start +// at 0, so this is driven on Pool directly -- a ring layer's read starts +// mid-pool, and dropping the start would silently return the wrong cells. +TEST_F(MLXSequenceCacheTest, PoolHonorsRunStart) { + using namespace ::mlx::core; + Pool p(/*slots=*/8, H, D, float16); + array x = randn(3, float16); + p.write(cache::Run{/*start=*/2, /*len=*/3}, x, s); + + EXPECT_TRUE(allclose(p.read(cache::Run{2, 3}, s), x, 0.0f)); + // The cells before the run are untouched, so reading from 0 is not the same + // window -- the regression this guards against. + EXPECT_FALSE(allclose(p.read(cache::Run{0, 3}, s), x, 0.0f)); +} + +} // namespace diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 3ab424830e0..40f2a28fdab 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -116,15 +116,17 @@ struct LayerConfig { }; // Model facts + runtime policy the byte layer sizes its pools from. capacity is -// the logical cap; initial_capacity tunes the byte layer's lazy-doubling pool; -// max_write is the max tokens written per step (a ring layer sizes its slots to -// window + max_write - 1 so a multi-token step fits); unset means each ring -// layer uses its own window. `layers` is per-layer: size 1 == uniform across -// all layers, else == n_layers. +// the logical cap; kv_dtype is the ET ScalarType the byte layer stores K/V in; +// initial_capacity tunes the byte layer's lazy-doubling pool; max_write is the +// max tokens written per step (a ring layer sizes its slots to window + +// max_write - 1 so a multi-token step fits); unset means each ring layer uses +// its own window. `layers` is per-layer: size 1 == uniform across all layers, +// else == n_layers. struct CacheConfig { int capacity; int n_layers; std::vector layers; + int kv_dtype; int initial_capacity = 512; std::optional max_write; };