Add off-graph KV-cache ring policy

extension/llm/cache/cache.h

@@ -13,9 +13,10 @@
 // as_control()/as_planner() (static upcasts -- no dynamic_cast/RTTI, no
 // diamond): a runner-facing control face (SequenceControl) and a backend-facing
 // planner face (SequencePlanner). One controller (SequenceCache) drives all
-// layers and dispatches per-layer layout to a LayoutPolicy (full-history flat
-// today), so a multi-layer model shares one logical length. Errors are plain
-// C++ (bool / std::optional); cache_et.h adapts them to Error/Result for ET
+// layers and dispatches per-layer layout to a LayoutPolicy (flat = full
+// history, ring = sliding window), so a mixed model (e.g. gemma4's alternating
+// full/sliding-window layers) shares one logical length. Errors are plain C++
+// (bool / std::optional); cache_et.h adapts them to Error/Result for ET
 // consumers.
 
 #include <optional>
@@ -58,10 +59,10 @@ struct Run {
 };
 
 // Integer-only handoff from the planner to the backend byte layer. Runs are in
-// logical order (oldest -> newest); a flat layer uses one run, a windowing
-// layer up to two (a write/read that wraps its buffer splits in two).
-// read_base_pos is the logical position of read[0].start (0 for flat), so the
-// backend can align RoPE / the attention mask.
+// logical order (oldest -> newest); a flat layer uses one run, a ring layer up
+// to two (a write/read that wraps the buffer splits in two). read_base_pos is
+// the logical position of read[0].start (0 for flat; the window start for
+// ring), so the backend can align RoPE / the attention mask.
 struct SeqStepPlan {
   Run write[2];
   int n_write;
@@ -84,26 +85,27 @@ class SequencePlanner {
   virtual void commit(const SeqStepPlan& plan) = 0;
 };
 
-// Per-layer layout behavior (e.g. full-history flat). Pure: plan() has no side
-// effects, so the controller (SequenceCache) owns length.
+// Per-layer layout behavior (flat = full history; ring = sliding window). Pure:
+// plan() has no side effects, so the controller (SequenceCache) owns length.
 class LayoutPolicy {
  public:
   virtual ~LayoutPolicy() = default;
   // Write/read runs for T cells at logical `position`. Precondition: T fits the
   // policy's window (the runner chunks prefill so a step fits).
   virtual SeqStepPlan plan(int position, int T) const = 0;
-  // Oldest logical position still retained given the current length (0 for a
-  // full-history policy; a windowing policy retains only its last window). Used
-  // to bound rewind.
+  // Oldest logical position still retained given the current length (0 for
+  // flat; length - window for ring). Used to bound rewind.
   virtual int retained_from(int length) const = 0;
 };
 
 // Per-layer cache kind and its parameters.
 struct LayerPolicy {
-  enum class Kind : int { Flat = 0 }; // serialized values: append-only
+  enum class Kind : int {
+    Flat = 0,
+    Ring = 1
+  }; // serialized values: append-only
   Kind kind = Kind::Flat;
-  int window =
-      0; // sliding-window size for windowing policies; 0 = full history
+  int window = 0; // Ring: sliding-window size; must be 0 when Flat
 };
 
 // Per-layer architecture facts + cache policy.
@@ -114,13 +116,17 @@ 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.
-// `layers` is per-layer: size 1 == uniform across all layers, else == n_layers.
+// 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.
 struct CacheConfig {
   int capacity;
   int n_layers;
   std::vector<LayerConfig> layers;
   int initial_capacity = 512;
+  std::optional<int> max_write;
 };
 
 } // namespace cache

extension/llm/cache/sequence_cache.h

@@ -8,10 +8,10 @@
 
 #pragma once
 
-// The neutral single-sequence controller (SequenceCache) and its flat layout
-// policy (FlatPolicy). SequenceCache owns the one logical length for the whole
-// model and dispatches per-layer layout to a policy, so a multi-layer model
-// stays coherent. Tensor-free / ET-independent.
+// The neutral single-sequence controller (SequenceCache) and its layout
+// policies (FlatPolicy / RingPolicy). SequenceCache owns the one logical length
+// for the whole model and dispatches per-layer layout to a policy, so a mixed
+// flat/ring model (gemma4) stays coherent. Tensor-free / ET-independent.
 
 #include <algorithm>
 #include <memory>
@@ -43,14 +43,60 @@ class FlatPolicy final : public LayoutPolicy {
   }
 };
 
+// Sliding window of `window` cells over a ring of `window + max_write - 1`
+// slots. The ring is oversized so a step of up to max_write tokens fits without
+// overwriting cells earlier queries in the same step still attend to; the
+// backend masks each query to its own window within the read span.
+class RingPolicy final : public LayoutPolicy {
+ public:
+  RingPolicy(int window, int max_write)
+      : window_(window), ring_size_(window + max_write - 1) {}
+
+  int retained_from(int length) const override {
+    return length > window_ ? length - window_ : 0;
+  }
+  SeqStepPlan plan(int position, int T) const override {
+    const int end = position + T;
+    SeqStepPlan p{};
+    // Write this step's T cells (T <= max_write), wrapping the ring if needed.
+    p.n_write = split_runs(position, T, p.write);
+    // Read the union of the step's per-query windows, [position - window + 1,
+    // end), in logical (oldest -> newest) order.
+    const int rstart = position - window_ + 1 > 0 ? position - window_ + 1 : 0;
+    const int rlen = end - rstart;
+    p.n_read = split_runs(rstart, rlen, p.read);
+    p.read_base_pos = rstart;
+    return p;
+  }
+
+ private:
+  // Split a logical run [start, start+len) into up to two physical runs in the
+  // ring (wrapping at ring_size_). Precondition: len <= ring_size_.
+  int split_runs(int start, int len, Run out[2]) const {
+    const int phys_start = start % ring_size_;
+    const int first_len = std::min(len, ring_size_ - phys_start);
+    out[0] = Run{phys_start, first_len};
+    if (first_len < len) {
+      out[1] = Run{0, len - first_len};
+      return 2;
+    }
+    return 1;
+  }
+
+  int window_;
+  int ring_size_;
+};
+
 // One controller for all layers: owns the single logical length, admission, and
 // rewind; dispatches per-layer layout to a shared LayoutPolicy. Policies are
-// deduped by (kind, window), so a uniform model holds a single policy object.
+// deduped by (kind, window), so a uniform or two-kind (gemma4) model holds one
+// or two policy objects.
 class SequenceCache : public CacheBase,
                       public SequenceControl,
                       public SequencePlanner {
  public:
-  explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity) {
+  explicit SequenceCache(const CacheConfig& cfg)
+      : capacity_(cfg.capacity), max_write_(cfg.max_write) {
     layer_to_policy_.reserve(cfg.n_layers);
     for (int l = 0; l < cfg.n_layers; ++l) {
       // layers size 1 = one config broadcast to every layer, else per-layer.
@@ -105,6 +151,13 @@ class SequenceCache : public CacheBase,
     if (position + T > capacity_) {
       return std::nullopt;
     }
+    // A ring layer's slots are sized window + max_write - 1 (max_write defaults
+    // to the window), so a step larger than that would overrun it.
+    const LayerPolicy& spec = specs_[layer_to_policy_[layer]];
+    if (spec.kind == LayerPolicy::Kind::Ring &&
+        T > (max_write_ ? *max_write_ : spec.window)) {
+      return std::nullopt;
+    }
     return policies_[layer_to_policy_[layer]]->plan(position, T);
   }
   void commit(const SeqStepPlan& plan) override {
@@ -129,11 +182,17 @@ class SequenceCache : public CacheBase,
     policies_.push_back(make_policy(lp));
     return static_cast<int>(policies_.size() - 1);
   }
-  std::unique_ptr<LayoutPolicy> make_policy(const LayerPolicy&) const {
+  std::unique_ptr<LayoutPolicy> make_policy(const LayerPolicy& lp) const {
+    if (lp.kind == LayerPolicy::Kind::Ring) {
+      // Unset max_write -> default a ring layer to its own window (ring 2w-1).
+      const int mw = max_write_ ? *max_write_ : lp.window;
+      return std::make_unique<RingPolicy>(lp.window, mw);
+    }
     return std::make_unique<FlatPolicy>();
   }
 
   int capacity_;
+  std::optional<int> max_write_;
   int length_ = 0;
   std::vector<LayerPolicy> specs_; // parallel to policies_, for dedup
   std::vector<std::unique_ptr<LayoutPolicy>> policies_;

extension/llm/cache/test/cache_test.cpp

@@ -33,6 +33,9 @@ namespace {
 LayerConfig flat_layer() {
   return LayerConfig{LayerPolicy{LayerPolicy::Kind::Flat, 0}, 2, 8};
 }
+LayerConfig ring_layer(int window) {
+  return LayerConfig{LayerPolicy{LayerPolicy::Kind::Ring, window}, 2, 8};
+}
 } // namespace
 
 // Initializes the ExecuTorch PAL so the ET adapter's error paths (which ET_LOG)
@@ -66,6 +69,58 @@ TEST_F(CacheTest, FlatPlanAppendsAndReadsAllHistory) {
   EXPECT_EQ(p1->read[0].len, 5);
 }
 
+// ---- Ring policy: wrap, eviction, read base position -----------------------
+
+TEST_F(CacheTest, RingPlanWrapsAndEvicts) {
+  // window 4, max_write unset -> defaults to window -> ring of 2*4 - 1 = 7.
+  SequenceCache cache(CacheConfig{100, 1, {ring_layer(4)}});
+
+  // Prefill chunk of 4 at position 0: one write run, reads [0,4), no wrap.
+  auto p0 = cache.plan(0, 0, 4);
+  ASSERT_TRUE(p0.has_value());
+  EXPECT_EQ(p0->n_write, 1);
+  EXPECT_EQ(p0->write[0].start, 0);
+  EXPECT_EQ(p0->write[0].len, 4);
+  EXPECT_EQ(p0->n_read, 1);
+  EXPECT_EQ(p0->read[0].start, 0);
+  EXPECT_EQ(p0->read[0].len, 4);
+  EXPECT_EQ(p0->read_base_pos, 0);
+  cache.commit(*p0);
+
+  // Decode at position 7: write wraps to slot 0 (7 % 7); the read window
+  // [4,8) wraps the ring -> phys [4,6] then [0,0]. Positions 0-3 evicted.
+  auto p1 = cache.plan(0, 7, 1);
+  ASSERT_TRUE(p1.has_value());
+  EXPECT_EQ(p1->n_write, 1);
+  EXPECT_EQ(p1->write[0].start, 0);
+  EXPECT_EQ(p1->write[0].len, 1);
+  EXPECT_EQ(p1->n_read, 2);
+  EXPECT_EQ(p1->read[0].start, 4);
+  EXPECT_EQ(p1->read[0].len, 3);
+  EXPECT_EQ(p1->read[1].start, 0);
+  EXPECT_EQ(p1->read[1].len, 1);
+  EXPECT_EQ(p1->read_base_pos, 4); // oldest retained is logical position 4
+
+  // A step larger than max_write would overrun the ring -> rejected.
+  EXPECT_FALSE(cache.plan(0, 0, 5).has_value());
+}
+
+// ---- Mixed flat/ring: one shared length across layers ----------------------
+
+TEST_F(CacheTest, MixedFlatRingShareOneLength) {
+  SequenceCache cache(CacheConfig{100, 2, {flat_layer(), ring_layer(4)}});
+
+  // Same step drives both layers (T <= window); commit once.
+  auto p = cache.plan(0, 0, 3);
+  ASSERT_TRUE(p.has_value());
+  ASSERT_TRUE(cache.plan(1, 0, 3).has_value());
+  cache.commit(*p);
+
+  // Length advanced once (not per layer): length == 3.
+  EXPECT_TRUE(cache.can_extend(97));
+  EXPECT_FALSE(cache.can_extend(98));
+}
+
 // ---- Admission / rewind ----------------------------------------------------
 
 TEST_F(CacheTest, CanExtendBoundedByCapacity) {
@@ -89,6 +144,21 @@ TEST_F(CacheTest, FlatRewindsFreelyToZero) {
   EXPECT_TRUE(cache.can_extend(8));
 }
 
+TEST_F(CacheTest, RewindBoundedByRingWindow) {
+  SequenceCache cache(CacheConfig{100, 2, {flat_layer(), ring_layer(4)}});
+  // Advance length to 10 in chunks of 2 (<= window).
+  for (int pos = 0; pos < 10; pos += 2) {
+    auto p = cache.plan(0, pos, 2);
+    ASSERT_TRUE(p.has_value());
+    ASSERT_TRUE(cache.plan(1, pos, 2).has_value());
+    cache.commit(*p);
+  }
+  // ring(4) has evicted everything older than 10 - 4 = 6.
+  EXPECT_FALSE(cache.rewind(5)); // older than the ring retains
+  EXPECT_TRUE(cache.rewind(6)); // exactly the floor
+  EXPECT_FALSE(cache.rewind(11)); // cannot grow
+}
+
 // ---- Faces / registry / session --------------------------------------------
 
 TEST_F(CacheTest, FaceRecoveryReturnsSameObject) {