[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_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::"

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 <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
+};
+
+// 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

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)

backends/mlx/runtime/MLXInterpreter.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include "MLXCache.h"
 #include "MLXExecutor.h"
 
 #include <mlx/array.h>
@@ -293,6 +294,74 @@ 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.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");
+  }
+  // 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<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);
+  // 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<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 +2028,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/runtime/MLXSequenceCache.h

@@ -0,0 +1,175 @@
+/*
+ * 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 "MLXCache.h" // AttendSpec, MLXCache
+#include "MLXExecutor.h" // resolve_dtype
+
+#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, 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<int>(buf_.shape(1));
+    const int D = static_cast<int>(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<int>(update.shape(2)) != run.len) {
+      throw std::runtime_error("Pool::write: update length != run length");
+    }
+    if (static_cast<int>(update.shape(1)) != H ||
+        static_cast<int>(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<int>(buf_.shape(1));
+    const int D = static_cast<int>(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<int>(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(checked(cfg)) {
+    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];
+      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<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
+
+    std::optional<cache::SeqStepPlan> 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<size_t>(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:
+  // 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<Pool> kpool_;
+  std::vector<Pool> vpool_;
+};
+
+} // namespace mlx
+} // namespace backends
+} // namespace executorch

backends/mlx/serialization/schema.fbs

@@ -153,6 +153,19 @@ table SdpaNode {
     causal: bool = false;
 }
 
+table UpdateAndAttendNode {
+    q: Tid (required);
+    k: Tid (required);
+    v: Tid (required);
+    position: Tid (required);
+    out: Tid (required);
+    // 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
+}
+
 table AddNode {
     a: Tid (required);
     b: Tid (required);
@@ -1170,7 +1183,8 @@ union OpNode {
     BitwiseOrNode,
     BitwiseXorNode,
     IfNode,
-    RandomBitsNode
+    RandomBitsNode,
+    UpdateAndAttendNode
     // BC: Add new op nodes here (append only)
 }