[MLX] Add off-graph KV cache ring runtime

backends/mlx/runtime/MLXCache.h

@@ -16,11 +16,10 @@ 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).
+// The K/V window to attend over + how to mask it; the cache owns the semantic
+// and hands over whatever MLX SDPA needs to apply it. `kind` mirrors MLX's mask
+// forms: no mask, its fused "causal", or an explicit tensor for anything MLX
+// cannot express -- a sliding window, and later tree/speculative patterns.
 struct AttendSpec {
   Tensor K;
   Tensor V;

backends/mlx/runtime/MLXSequenceCache.h

@@ -117,15 +117,33 @@ class Pool {
   Tensor buf_;
 };
 
+// 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 Tensor window_causal_mask(int T, int S, int window, StreamOrDevice s) {
+  using namespace ::mlx::core;
+  Tensor 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;
+  Tensor 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);
+}
+
 // 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.
+// 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 +152,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 +186,80 @@ 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.
+    if (T == 1) {
+      return AttendSpec{K, V, AttendSpec::Mask::None, std::nullopt};
+    }
+    // A ring layer bounds each query to its own window. That only bites when
+    // the window is narrower than the span (a multi-token step reads the union
+    // of its queries' windows, window + T - 1 cells); otherwise plain causal is
+    // exact and MLX applies it fused, with no mask tensor.
+    const int S = static_cast<int>(K.shape(2));
+    if (window_[l] > 0 && window_[l] < S) {
+      return AttendSpec{
+          K,
+          V,
+          AttendSpec::Mask::Explicit,
+          window_causal_mask(T, S, window_[l], s)};
+    }
+    // MLX "causal" is lower-right aligned, so fresh and chunked prefill are
+    // both correct with the new tokens at the tail.
+    return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt};
   }
 
  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 +272,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

@@ -63,6 +63,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 +276,137 @@ 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(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(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 stays Causal while its window still covers the span, so MLX
+// applies its fused mask and no tensor is built. A flat layer always does.
+TEST_F(MLXSequenceCacheTest, WindowWiderThanSpanStaysCausal) {
+  using namespace ::mlx::core;
+  MLXSequenceCache ring(ring_config(
+      /*capacity=*/64,
+      /*window=*/4,
+      /*max_write=*/2,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half)));
+  array k = randn(2, float16); // span 2 < window 4
+  AttendSpec rspec = ring.update_and_fetch(0, /*position=*/0, k, k, s);
+  EXPECT_EQ(rspec.kind, AttendSpec::Mask::Causal);
+  EXPECT_FALSE(rspec.mask.has_value());
+
+  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.mask.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. Since the window is now narrower than the span, the cache hands back a
+  // band to narrow each query back to its own 4.
+  EXPECT_EQ(spec.K.shape(2), W + 1);
+  EXPECT_EQ(spec.kind, AttendSpec::Mask::Explicit);
+  ASSERT_TRUE(spec.mask.has_value());
+  // Span indices 0..4 are positions 1..5. Query 0 is position 4 and attends
+  // 1..4; query 1 is position 5 and attends 2..5 -- four keys each.
+  // clang-format off
+  std::vector<int> band = {
+      1, 1, 1, 1, 0,
+      0, 1, 1, 1, 1};
+  // clang-format on
+  EXPECT_TRUE(allclose(
+      astype(*spec.mask, int32, s),
+      array(band.data(), Shape{1, 1, 2, W + 1}, int32),
+      0.0f));
+  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) {