[MLX] Add off-graph KV cache ring runtime

backends/mlx/runtime/MLXCache.h

@@ -17,15 +17,19 @@ 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).
+// `kind` mirrors MLX SDPA's mask forms (no mask / "causal" / explicit tensor),
+// with `window` as the declarative sliding-window modifier (the mask_mod axis).
+// The cache states the semantic; the handler materializes whatever the backend
+// cannot express natively. Will be added: 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
+  // Causal only: each query attends at most this many keys (itself included).
+  // Unset = unbounded history.
+  std::optional<int> window;
 };
 
 // Tensor-typed op face of the off-graph KV cache, kept separate from the

backends/mlx/runtime/MLXInterpreter.h

@@ -294,6 +294,24 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) {
   st.set_tensor(n.out, std::move(out));
 }
 
+// Bool mask [1, 1, T, S] for T queries over a span of S keys, where each query
+// attends at most `window` keys ending at itself. The span is right-aligned
+// (the newest key belongs to the last query), so query i spans keys
+// j - i <= S - T -- the same bound MLX's "causal" applies -- and the window
+// adds the lower bound j - i > S - T - window.
+inline array window_causal_mask(int T, int S, int window, StreamOrDevice s) {
+  array diff = subtract(
+      reshape(arange(S, int32, s), Shape{1, S}, s),
+      reshape(arange(T, int32, s), Shape{T, 1}, s),
+      s);
+  const int hi = S - T;
+  array band = logical_and(
+      less_equal(diff, array(hi), s),
+      greater_equal(diff, array(hi - window + 1), s),
+      s);
+  return reshape(band, Shape{1, 1, T, S}, s);
+}
+
 inline void exec_update_and_attend(
     const UpdateAndAttendNode& n,
     ExecutionState& st,
@@ -333,12 +351,22 @@ inline void exec_update_and_attend(
   // only by spec.mask, so an Explicit with no mask would silently attend
   // unmasked.
   std::string mask_mode;
+  std::optional<array> mask = spec.mask;
   switch (spec.kind) {
     case AttendSpec::Mask::None:
       break;
-    case AttendSpec::Mask::Causal:
-      mask_mode = "causal";
+    case AttendSpec::Mask::Causal: {
+      const int T = static_cast<int>(q.shape(2));
+      const int S = static_cast<int>(K.shape(2));
+      // MLX has no sliding-window mode, so a window that actually bites has to
+      // be materialized. It never bites when it covers the whole span.
+      if (spec.window && *spec.window < S) {
+        mask = window_causal_mask(T, S, *spec.window, s);
+      } else {
+        mask_mode = "causal";
+      }
       break;
+    }
     case AttendSpec::Mask::Explicit:
       if (!spec.mask) {
         throw std::runtime_error(
@@ -347,14 +375,7 @@ inline void exec_update_and_attend(
       break;
   }
   array out = fast::scaled_dot_product_attention(
-      q,
-      K,
-      V,
-      static_cast<float>(*n.scale),
-      mask_mode,
-      spec.mask,
-      std::nullopt,
-      s);
+      q, K, V, static_cast<float>(*n.scale), mask_mode, 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);

backends/mlx/runtime/MLXSequenceCache.h

@@ -122,10 +122,9 @@ class Pool {
 // 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.
+// Flat and ring layers can be mixed (gemma4 alternates them): the planner emits
+// physical runs either way, so they differ here only in how many slots the pool
+// holds and whether a step's runs wrap.
 class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
  public:
   explicit MLXSequenceCache(const cache::CacheConfig& cfg)
@@ -134,21 +133,21 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
         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));
+    window_.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 may reach the full cap. A
-      // ring layer will instead cap at window + max_write - 1 slots.
-      const int max_slots = cfg.capacity;
-      kpool_.emplace_back(
-          cfg.initial_capacity, max_slots, lc.n_kv_heads, lc.head_dim, dt);
-      vpool_.emplace_back(
-          cfg.initial_capacity, max_slots, lc.n_kv_heads, lc.head_dim, dt);
+      const bool ring = lc.policy.kind == cache::LayerPolicy::Kind::Ring;
+      window_.push_back(ring ? lc.policy.window : 0);
+      // Flat retains all history, so its pool may reach the full cap and starts
+      // small. A ring layer recycles a fixed window + max_write - 1 slots.
+      const int max_slots = ring ? lc.policy.window +
+              (cfg.max_write ? *cfg.max_write : lc.policy.window) - 1
+                                 : cfg.capacity;
+      const int initial = ring ? max_slots : cfg.initial_capacity;
+      kpool_.emplace_back(initial, max_slots, lc.n_kv_heads, lc.head_dim, dt);
+      vpool_.emplace_back(initial, max_slots, lc.n_kv_heads, lc.head_dim, dt);
     }
   }
 
@@ -168,27 +167,75 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
       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);
+    write_runs(kpool_[l], p->write, p->n_write, k, s);
+    write_runs(vpool_[l], p->write, p->n_write, v, s);
+    Tensor K = read_runs(kpool_[l], p->read, p->n_read, s);
+    Tensor V = read_runs(vpool_[l], p->read, p->n_read, 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};
+    // a single decode token needs no mask -- its span is exactly what it may
+    // attend, on a ring layer as much as a flat one. A ring layer additionally
+    // bounds each query to its own window; the handler materializes that only
+    // when the window is narrower than the span.
+    if (T == 1) {
+      return AttendSpec{
+          K, V, AttendSpec::Mask::None, std::nullopt, std::nullopt};
+    }
+    std::optional<int> window;
+    if (window_[l] > 0) {
+      window = window_[l];
+    }
+    return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt, window};
   }
 
  private:
+  // Scatter `update` across the step's runs. Runs are in logical order, so
+  // consecutive slices of `update` map to consecutive runs. A flat step is one
+  // run; a ring step splits in two when it wraps the pool.
+  static void write_runs(
+      Pool& pool,
+      const cache::Run* runs,
+      int n,
+      const Tensor& update,
+      StreamOrDevice s) {
+    if (n == 1) {
+      pool.write(runs[0], update, s);
+      return;
+    }
+    const int H = static_cast<int>(update.shape(1));
+    const int D = static_cast<int>(update.shape(3));
+    int off = 0;
+    for (int i = 0; i < n; ++i) {
+      pool.write(
+          runs[i],
+          ::mlx::core::slice(
+              update,
+              ::mlx::core::Shape{0, 0, off, 0},
+              ::mlx::core::Shape{1, H, off + runs[i].len, D},
+              ::mlx::core::Shape{1, 1, 1, 1},
+              s),
+          s);
+      off += runs[i].len;
+    }
+  }
+
+  // Gather the step's runs into the retained window, oldest -> newest.
+  static Tensor
+  read_runs(const Pool& pool, const cache::Run* runs, int n, StreamOrDevice s) {
+    if (n == 1) {
+      return pool.read(runs[0], s);
+    }
+    std::vector<Tensor> parts;
+    parts.reserve(static_cast<size_t>(n));
+    for (int i = 0; i < n; ++i) {
+      parts.push_back(pool.read(runs[i], s));
+    }
+    return ::mlx::core::concatenate(parts, 2, s);
+  }
+
   // 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.
@@ -201,6 +248,7 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
 
   std::vector<Pool> kpool_;
   std::vector<Pool> vpool_;
+  std::vector<int> window_; // per layer; 0 = flat (unbounded history)
 };
 
 } // namespace mlx

backends/mlx/test/mlx_sequence_cache_test.cpp

@@ -17,6 +17,7 @@
 //
 // Must run on Apple Silicon: MLX needs the Metal backend.
 
+#include "MLXInterpreter.h" // window_causal_mask
 #include "MLXSequenceCache.h"
 
 #include <mlx/mlx.h>
@@ -63,6 +64,22 @@ cache::CacheConfig flat_config(
   return cfg;
 }
 
+// One ring layer of `window` cells; max_write bounds a multi-token step.
+cache::CacheConfig ring_config(
+    int capacity,
+    int window,
+    int max_write,
+    int n_kv_heads,
+    int head_dim,
+    int kv_dtype) {
+  cache::CacheConfig cfg =
+      flat_config(capacity, /*n_layers=*/1, n_kv_heads, head_dim, kv_dtype);
+  cfg.layers[0].policy =
+      cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, window};
+  cfg.max_write = max_write;
+  return cfg;
+}
+
 class MLXSequenceCacheTest : public ::testing::Test {
  protected:
   const int H = 2;
@@ -260,6 +277,127 @@ TEST_F(MLXSequenceCacheTest, ZeroInitialCapacityGrowsOnFirstWrite) {
   EXPECT_TRUE(allclose(spec0.K, k0, 0.0f));
 }
 
+// The sliding-window mask is a band on (key - query): causal above, window
+// below. The span is right-aligned, so the last query's own key is the last
+// key, and each query attends `window` keys ending at its own.
+TEST_F(MLXSequenceCacheTest, WindowCausalMaskIsABand) {
+  using namespace ::mlx::core;
+  // T=3 queries over S=5 keys, window 3. Query i owns key i + (S - T), and
+  // attends the `window` keys ending there -- a band, one row per query.
+  // clang-format off
+  std::vector<int> want = {
+      1, 1, 1, 0, 0,
+      0, 1, 1, 1, 0,
+      0, 0, 1, 1, 1};
+  // A window covering the whole span leaves plain causal: no lower bound bites.
+  std::vector<int> causal = {
+      1, 1, 1, 0, 0,
+      1, 1, 1, 1, 0,
+      1, 1, 1, 1, 1};
+  // clang-format on
+  array m = astype(ops::window_causal_mask(3, 5, 3, s), int32, s);
+  EXPECT_TRUE(allclose(m, array(want.data(), Shape{1, 1, 3, 5}, int32), 0.0f));
+
+  array full = astype(ops::window_causal_mask(3, 5, 5, s), int32, s);
+  EXPECT_TRUE(
+      allclose(full, array(causal.data(), Shape{1, 1, 3, 5}, int32), 0.0f));
+}
+
+// A ring layer decodes past its window: the read span stops at `window` cells
+// and holds the newest tokens, evicting the oldest. No mask -- a single query
+// may attend its whole span.
+TEST_F(MLXSequenceCacheTest, RingDecodeEvictsOldestAndNeedsNoMask) {
+  using namespace ::mlx::core;
+  const int W = 4;
+  MLXSequenceCache c(ring_config(
+      /*capacity=*/64,
+      /*window=*/W,
+      /*max_write=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half)));
+
+  // Feed 6 single tokens; the last 4 must survive, oldest -> newest.
+  std::vector<array> toks;
+  for (int i = 0; i < 6; ++i) {
+    toks.push_back(randn(1, float16));
+  }
+  AttendSpec spec{toks[0], toks[0], AttendSpec::Mask::None, {}, {}};
+  for (int i = 0; i < 6; ++i) {
+    spec = c.update_and_fetch(0, /*position=*/i, toks[i], toks[i], s);
+  }
+  EXPECT_EQ(spec.kind, AttendSpec::Mask::None);
+  EXPECT_EQ(spec.K.shape(2), W);
+  EXPECT_TRUE(allclose(
+      spec.K,
+      concatenate(std::vector<array>{toks[2], toks[3], toks[4], toks[5]}, 2, s),
+      0.0f));
+}
+
+// A ring layer's multi-token step declares its window; the handler decides
+// whether to materialize it. Flat layers declare no window.
+TEST_F(MLXSequenceCacheTest, RingDeclaresWindowOnCausal) {
+  using namespace ::mlx::core;
+  const int W = 4;
+  MLXSequenceCache ring(ring_config(
+      /*capacity=*/64,
+      /*window=*/W,
+      /*max_write=*/2,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half)));
+  array k = randn(2, float16);
+  AttendSpec rspec = ring.update_and_fetch(0, /*position=*/0, k, k, s);
+  EXPECT_EQ(rspec.kind, AttendSpec::Mask::Causal);
+  ASSERT_TRUE(rspec.window.has_value());
+  EXPECT_EQ(*rspec.window, W);
+
+  MLXSequenceCache flat(flat_config(
+      /*capacity=*/64,
+      /*n_layers=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half)));
+  AttendSpec fspec = flat.update_and_fetch(0, /*position=*/0, k, k, s);
+  EXPECT_EQ(fspec.kind, AttendSpec::Mask::Causal);
+  EXPECT_FALSE(fspec.window.has_value());
+}
+
+// A step whose runs wrap the ring is scattered and gathered in logical order.
+TEST_F(MLXSequenceCacheTest, RingStepWrapsAndRejoinsInOrder) {
+  using namespace ::mlx::core;
+  const int W = 4;
+  const int MW = 2;
+  MLXSequenceCache c(ring_config(
+      /*capacity=*/64,
+      /*window=*/W,
+      /*max_write=*/MW,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half)));
+
+  // ring_size = W + MW - 1 = 5. Fill 4 tokens, then a 2-token step at
+  // position 4 writes slots 4 and 0 -- a wrap.
+  std::vector<array> toks;
+  for (int i = 0; i < 4; ++i) {
+    toks.push_back(randn(1, float16));
+    c.update_and_fetch(0, /*position=*/i, toks.back(), toks.back(), s);
+  }
+  array pair = randn(2, float16);
+  AttendSpec spec = c.update_and_fetch(0, /*position=*/4, pair, pair, s);
+
+  // The span is the union of the two queries' windows -- position 4 attends
+  // 1..4 and position 5 attends 2..5 -- so it is window + T - 1 = 5 cells, not
+  // 4. Narrowing each query back to its own 4 is the declared window's job.
+  EXPECT_EQ(spec.K.shape(2), W + 1);
+  ASSERT_TRUE(spec.window.has_value());
+  EXPECT_EQ(*spec.window, W);
+  EXPECT_TRUE(allclose(
+      spec.K,
+      concatenate(std::vector<array>{toks[1], toks[2], toks[3], pair}, 2, s),
+      0.0f));
+}
+
 // A negative initial_capacity is rejected rather than reaching MLX as a
 // negative dimension.
 TEST_F(MLXSequenceCacheTest, NegativeInitialCapacityThrows) {