[MLX] Add off-graph KV-cache flat runtime

.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::"

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 <optional>
+
+#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<Tensor> 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

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)

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 <optional>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "MLXCacheImpl.h" // AttendSpec, MLXCacheImpl
+#include "MLXExecutor.h" // to_shape
+
+#include <executorch/extension/llm/cache/cache.h>
+#include <executorch/extension/llm/cache/sequence_cache.h>
+
+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<int>{1, H, capacity, D}),
+            dtype)) {}
+
+  void write(int start, const Tensor& update, StreamOrDevice s) {
+    const int cap = static_cast<int>(buf_.shape(2));
+    const int H = static_cast<int>(buf_.shape(1));
+    const int D = static_cast<int>(buf_.shape(3));
+    const int T = static_cast<int>(update.shape(2));
+    if (start < 0 || start + T > cap) {
+      throw std::runtime_error("FlatPool::write: run out of bounds");
+    }
+    if (static_cast<int>(update.shape(1)) != H ||
+        static_cast<int>(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<int> lo{0, 0, start, 0};
+    std::vector<int> 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<int>(buf_.shape(1));
+    const int D = static_cast<int>(buf_.shape(3));
+    std::vector<int> lo{0, 0, 0, 0};
+    std::vector<int> hi{1, H, len, D};
+    std::vector<int> 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<int8_t>(*cfg.kv_dtype));
+    kpool_.reserve(static_cast<size_t>(cfg.n_layers));
+    vpool_.reserve(static_cast<size_t>(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<int>(kpool_.size())) {
+      throw std::out_of_range("update_and_fetch: layer out of range");
+    }
+    const int T = static_cast<int>(k.shape(2)); // BHSD: seq axis is 2
+    const int start = position;
+
+    std::optional<cache::SeqStepPlan> 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<size_t>(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<FlatPool> kpool_;
+  std::vector<FlatPool> vpool_;
+};
+
+} // namespace mlx
+} // namespace backends
+} // namespace executorch

backends/mlx/runtime/MLXInterpreter.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "MLXCacheImpl.h"
 #include "MLXExecutor.h"
 
 #include <mlx/array.h>
@@ -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<int32_t>()[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<float>(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<SdpaNode>(instr.node), st, s);
         break;
+      case OpCode::UPDATE_AND_ATTEND:
+        ops::exec_update_and_attend(
+            std::get<UpdateAndAttendNode>(instr.node), st, s);
+        break;
       case OpCode::ADD:
         ops::exec_add(std::get<AddNode>(instr.node), st, s);
         break;

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)
 }
 

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()