From 011bd321a55c4fbc9187601538862696c69faf4f Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Thu, 30 Jul 2026 11:51:04 -0700 Subject: [PATCH 1/4] [MLX] Add off-graph KV-cache flat runtime --- .github/workflows/mlx.yml | 6 +- backends/mlx/runtime/MLXCacheImpl.h | 52 ++++++ backends/mlx/runtime/MLXExecutor.h | 4 + backends/mlx/runtime/MLXFlatCache.h | 146 +++++++++++++++++ backends/mlx/runtime/MLXInterpreter.h | 49 ++++++ backends/mlx/serialization/schema.fbs | 14 +- backends/mlx/test/CMakeLists.txt | 24 +++ backends/mlx/test/mlx_flat_cache_test.cpp | 184 ++++++++++++++++++++++ extension/llm/cache/cache.h | 3 + 9 files changed, 480 insertions(+), 2 deletions(-) create mode 100644 backends/mlx/runtime/MLXCacheImpl.h create mode 100644 backends/mlx/runtime/MLXFlatCache.h create mode 100644 backends/mlx/test/mlx_flat_cache_test.cpp 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 From 9c609a62f7316677c9e6de0eaa4d072cada838e4 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 31 Jul 2026 10:17:45 -0700 Subject: [PATCH 2/4] [MLX] Reshape the KV cache around per-layer policy and require kv_dtype --- .github/workflows/mlx.yml | 6 +- .../runtime/{MLXCacheImpl.h => MLXCache.h} | 11 +- backends/mlx/runtime/MLXExecutor.h | 5 +- backends/mlx/runtime/MLXFlatCache.h | 146 ---------------- backends/mlx/runtime/MLXInterpreter.h | 2 +- backends/mlx/runtime/MLXSequenceCache.h | 165 ++++++++++++++++++ backends/mlx/test/CMakeLists.txt | 14 +- ...e_test.cpp => mlx_sequence_cache_test.cpp} | 99 +++++------ extension/llm/cache/cache.h | 15 +- 9 files changed, 232 insertions(+), 231 deletions(-) rename backends/mlx/runtime/{MLXCacheImpl.h => MLXCache.h} (76%) delete mode 100644 backends/mlx/runtime/MLXFlatCache.h create mode 100644 backends/mlx/runtime/MLXSequenceCache.h rename backends/mlx/test/{mlx_flat_cache_test.cpp => mlx_sequence_cache_test.cpp} (57%) diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 0ce0f3b57a7..800bc91419f 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -66,15 +66,15 @@ 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 mlx_flat_cache_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 flat KV-cache op test" - ./cmake-out/backends/mlx/test/mlx_flat_cache_test + 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" diff --git a/backends/mlx/runtime/MLXCacheImpl.h b/backends/mlx/runtime/MLXCache.h similarity index 76% rename from backends/mlx/runtime/MLXCacheImpl.h rename to backends/mlx/runtime/MLXCache.h index b6f160e5c4a..8c4e45d8cbe 100644 --- a/backends/mlx/runtime/MLXCacheImpl.h +++ b/backends/mlx/runtime/MLXCache.h @@ -28,11 +28,14 @@ struct AttendSpec { 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 { +// 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 ~MLXCacheImpl() = default; + 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 -- diff --git a/backends/mlx/runtime/MLXExecutor.h b/backends/mlx/runtime/MLXExecutor.h index 2c0bca4b27a..eea49e40eb5 100644 --- a/backends/mlx/runtime/MLXExecutor.h +++ b/backends/mlx/runtime/MLXExecutor.h @@ -148,14 +148,13 @@ struct MutableBufferData { } }; -class MLXCacheImpl; // off-graph KV cache op face (MLXCacheImpl.h) +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 - MLXCacheImpl* cache{ - nullptr}; // Per-session off-graph KV cache; survives reset() + 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/MLXFlatCache.h b/backends/mlx/runtime/MLXFlatCache.h deleted file mode 100644 index 52bedc8f5d6..00000000000 --- a/backends/mlx/runtime/MLXFlatCache.h +++ /dev/null @@ -1,146 +0,0 @@ -/* - * 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 5de29233a22..4de788ab056 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -8,7 +8,7 @@ #pragma once -#include "MLXCacheImpl.h" +#include "MLXCache.h" #include "MLXExecutor.h" #include diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h new file mode 100644 index 00000000000..4c397eebeaa --- /dev/null +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -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. + */ + +#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. + 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/test/CMakeLists.txt b/backends/mlx/test/CMakeLists.txt index 6fa0ad50b39..9b9d5cb4f85 100644 --- a/backends/mlx/test/CMakeLists.txt +++ b/backends/mlx/test/CMakeLists.txt @@ -89,26 +89,26 @@ 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 +# 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_flat_cache_test + mlx_sequence_cache_test SOURCES - ${CMAKE_CURRENT_LIST_DIR}/mlx_flat_cache_test.cpp + ${CMAKE_CURRENT_LIST_DIR}/mlx_sequence_cache_test.cpp EXTRA_LIBS mlxdelegate mlx_schema mlx ) target_include_directories( - mlx_flat_cache_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../runtime + mlx_sequence_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 + mlx_sequence_cache_test PRIVATE -fsanitize=address,undefined + -fno-omit-frame-pointer ) target_link_options( - mlx_flat_cache_test PRIVATE ${_mlx_sanitizer_link_options} + mlx_sequence_cache_test PRIVATE ${_mlx_sanitizer_link_options} ) endif() diff --git a/backends/mlx/test/mlx_flat_cache_test.cpp b/backends/mlx/test/mlx_sequence_cache_test.cpp similarity index 57% rename from backends/mlx/test/mlx_flat_cache_test.cpp rename to backends/mlx/test/mlx_sequence_cache_test.cpp index bc9ccb03ea6..34a5fa40d7a 100644 --- a/backends/mlx/test/mlx_flat_cache_test.cpp +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -6,26 +6,23 @@ * LICENSE file in the root directory of this source tree. */ -// Op-level test for the off-graph flat KV cache (MLXFlatCache / FlatPool). +// Op-level test for the off-graph KV cache (MLXSequenceCache / Pool). // -// 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. +// 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 SDPA needs the Metal backend. +// Must run on Apple Silicon: MLX needs the Metal backend. -#include "MLXFlatCache.h" +#include "MLXSequenceCache.h" #include #include -#include -#include -#include #include using namespace ::executorch::backends::mlx; @@ -49,7 +46,7 @@ cache::CacheConfig flat_config( int n_layers, int n_kv_heads, int head_dim, - std::optional kv_dtype = std::nullopt) { + int kv_dtype) { cache::CacheConfig cfg; cfg.capacity = capacity; cfg.n_layers = n_layers; @@ -61,11 +58,10 @@ cache::CacheConfig flat_config( return cfg; } -class MLXFlatCacheTest : public ::testing::Test { +class MLXSequenceCacheTest : 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) { @@ -75,41 +71,29 @@ class MLXFlatCacheTest : public ::testing::Test { }; // Prefill: T=4 at position 0 -> Causal (lower-right aligned). -TEST_F(MLXFlatCacheTest, PrefillIsCausal) { +TEST_F(MLXSequenceCacheTest, PrefillIsCausal) { using namespace ::mlx::core; - MLXFlatCache c(flat_config( + MLXSequenceCache 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)); + EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); + EXPECT_TRUE(allclose(spec0.V, v0, 0.0f)); } -// Decode: after a T=4 prefill, a single query at position 4 -> None and must -// attend the full assembled history. -TEST_F(MLXFlatCacheTest, DecodeAttendsFullHistory) { +// 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; - MLXFlatCache c(flat_config( + MLXSequenceCache c(flat_config( /*capacity=*/32, /*n_layers=*/1, H, @@ -120,32 +104,20 @@ TEST_F(MLXFlatCacheTest, DecodeAttendsFullHistory) { 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)); + 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(MLXFlatCacheTest, StepPastCapacityThrows) { +TEST_F(MLXSequenceCacheTest, StepPastCapacityThrows) { using namespace ::mlx::core; - MLXFlatCache c(flat_config( + MLXSequenceCache c(flat_config( /*capacity=*/32, /*n_layers=*/1, H, @@ -157,9 +129,9 @@ TEST_F(MLXFlatCacheTest, StepPastCapacityThrows) { // 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) { +TEST_F(MLXSequenceCacheTest, StorageDtypeDiffersCastsOnWrite) { using namespace ::mlx::core; - MLXFlatCache c16(flat_config( + MLXSequenceCache c16(flat_config( /*capacity=*/32, /*n_layers=*/1, H, @@ -175,10 +147,19 @@ TEST_F(MLXFlatCacheTest, StorageDtypeDiffersCastsOnWrite) { 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}); +// 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 c657571a490..40f2a28fdab 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -116,20 +116,19 @@ 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; - // ET ScalarType storage precision; a backend that stores K/V may require it - // (MLX raises if unset). - std::optional kv_dtype; }; } // namespace cache From f96535c5e0fef5332dba8a2d3eaaa793a8a3d813 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 31 Jul 2026 10:29:34 -0700 Subject: [PATCH 3/4] [MLX] Require an explicit scale and switch exhaustively on mask kind --- backends/mlx/runtime/MLXInterpreter.h | 33 ++++++++++++++++++++----- backends/mlx/runtime/MLXSequenceCache.h | 5 ++++ backends/mlx/serialization/schema.fbs | 3 ++- 3 files changed, 34 insertions(+), 7 deletions(-) diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 4de788ab056..e85ae894d80 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -301,13 +301,17 @@ inline void exec_update_and_attend( 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 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. + // 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]; @@ -321,12 +325,29 @@ inline void exec_update_and_attend( // 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" : ""; + // 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), + static_cast(*n.scale), mask_mode, spec.mask, std::nullopt, diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 4c397eebeaa..1c1fabe8054 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -37,6 +37,11 @@ class Pool { // 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)); diff --git a/backends/mlx/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs index 0f4dfe0687a..99be50292db 100644 --- a/backends/mlx/serialization/schema.fbs +++ b/backends/mlx/serialization/schema.fbs @@ -160,7 +160,8 @@ table UpdateAndAttendNode { position: Tid (required); out: Tid (required); layer_id: int32; - scale: float; + 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 } From 1e895e46835540d2f8879d83e3a511c86aed3688 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 31 Jul 2026 12:10:52 -0700 Subject: [PATCH 4/4] [MLX] Validate the cache config and require layer_id --- backends/mlx/runtime/MLXInterpreter.h | 5 ++++- backends/mlx/runtime/MLXSequenceCache.h | 17 +++++++++++------ backends/mlx/serialization/schema.fbs | 7 ++++--- backends/mlx/test/mlx_sequence_cache_test.cpp | 15 +++++++++++++-- extension/llm/cache/cache.h | 10 ++++++++++ 5 files changed, 42 insertions(+), 12 deletions(-) diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index e85ae894d80..51946e966a6 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -301,6 +301,9 @@ inline void exec_update_and_attend( if (!st.cache) { throw std::runtime_error("update_and_attend: no cache installed"); } + if (!n.layer_id) { + throw std::runtime_error("update_and_attend: layer_id is not set"); + } if (!n.scale) { throw std::runtime_error("update_and_attend: scale is not set"); } @@ -316,7 +319,7 @@ inline void exec_update_and_attend( eval(pos); const int position = pos.data()[0]; AttendSpec spec = st.cache->update_and_fetch( - n.layer_id, + *n.layer_id, position, st.const_tensor_ref(n.k), st.const_tensor_ref(n.v), diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 1c1fabe8054..3f1bb7c4639 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -37,11 +37,6 @@ class Pool { // 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)); @@ -103,7 +98,7 @@ class Pool { class MLXSequenceCache : public cache::SequenceCache, public MLXCache { public: explicit MLXSequenceCache(const cache::CacheConfig& cfg) - : cache::SequenceCache(cfg) { + : cache::SequenceCache(checked(cfg)) { const ::mlx::core::Dtype dt = resolve_dtype(static_cast(cfg.kv_dtype)); kpool_.reserve(static_cast(cfg.n_layers)); @@ -161,6 +156,16 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { } private: + // Enforce the neutral contract as an exception, the failure mode this layer + // already uses. Runs as the base initializer's argument because + // SequenceCache's own ctor indexes `layers` before this class's body does. + static const cache::CacheConfig& checked(const cache::CacheConfig& cfg) { + if (!cache::valid(cfg)) { + throw std::runtime_error("MLXSequenceCache: invalid CacheConfig"); + } + return cfg; + } + std::vector kpool_; std::vector vpool_; }; diff --git a/backends/mlx/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs index 99be50292db..c69f5f48dda 100644 --- a/backends/mlx/serialization/schema.fbs +++ b/backends/mlx/serialization/schema.fbs @@ -159,9 +159,10 @@ table UpdateAndAttendNode { 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 + // Required, but FlatBuffers cannot mark a scalar (required): null-defaulted + // so an omitted field is distinguishable from 0 and the handler rejects it. + layer_id: int32 = null; + scale: float = null; out_dtype: int8 = null; // ET ScalarType; null = keep SDPA's native output } diff --git a/backends/mlx/test/mlx_sequence_cache_test.cpp b/backends/mlx/test/mlx_sequence_cache_test.cpp index 34a5fa40d7a..9e38f44fb52 100644 --- a/backends/mlx/test/mlx_sequence_cache_test.cpp +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -65,8 +65,7 @@ class MLXSequenceCacheTest : public ::testing::Test { ::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); + return ::mlx::core::random::normal(::mlx::core::Shape{1, H, T, D}, dt); } }; @@ -162,4 +161,16 @@ TEST_F(MLXSequenceCacheTest, PoolHonorsRunStart) { EXPECT_FALSE(allclose(p.read(cache::Run{0, 3}, s), x, 0.0f)); } +// A partial per-layer list is rejected instead of indexing past the end. +TEST_F(MLXSequenceCacheTest, PartialLayerListThrows) { + cache::CacheConfig cfg = flat_config( + /*capacity=*/32, + /*n_layers=*/4, + H, + D, + static_cast(ScalarType::Half)); + cfg.layers.push_back(cfg.layers.front()); // size 2, neither 1 nor n_layers + EXPECT_ANY_THROW(MLXSequenceCache{cfg}); +} + } // namespace diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 40f2a28fdab..21982bed893 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -131,6 +131,16 @@ struct CacheConfig { std::optional max_write; }; +// Whether `cfg` satisfies the contract above. Callers must check this before +// constructing a cache: the `layers` broadcast rule is indexed directly, so a +// list that is neither size 1 nor n_layers reads past the end. Reported as a +// bool rather than thrown, so each backend picks its own failure mode. +inline bool valid(const CacheConfig& cfg) { + return cfg.capacity > 0 && cfg.n_layers > 0 && + (cfg.layers.size() == 1 || + cfg.layers.size() == static_cast(cfg.n_layers)); +} + } // namespace cache } // namespace llm } // namespace extension