Add off-graph KV-cache SequenceCache and flat policy

extension/llm/cache/cache.h

@@ -8,29 +8,112 @@
 
 #pragma once
 
-// Neutral, ET-independent ownership handle for an off-graph KV cache. The
-// registry owns a cache as a CacheBase* -- an opaque, cache-agnostic anchor --
-// and hands it back by key. The typed faces a runner and backend use to drive
-// the cache are added by the concrete cache implementation.
+// Neutral, tensor-free, ET-independent KV-cache core shared across backends. A
+// cache exposes two faces recovered from the owning CacheBase* via
+// 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
+// consumers.
+
+#include <optional>
+#include <vector>
 
 namespace executorch {
 namespace extension {
 namespace llm {
 namespace cache {
 
-// Ownership / erasure anchor: the registry owns and deletes a cache through
-// this base, staying agnostic to the concrete cache type.
+class SequenceControl;
+class SequencePlanner;
+
+// Registry ownership / erasure anchor. The registry owns a cache as a
+// CacheBase*; the runner recovers the control face and the backend recovers the
+// planner face through these accessors (each concrete cache returns `this`).
 class CacheBase {
  public:
   virtual ~CacheBase() = default;
+  virtual SequenceControl* as_control() = 0;
+  virtual SequencePlanner* as_planner() = 0;
+};
+
+// Application (runner) face: lifecycle + admission, tensor-free.
+class SequenceControl {
+ public:
+  virtual ~SequenceControl() = default;
+  virtual bool can_extend(int n = 1) const = 0; // admission / hard-stop
+  virtual int capacity() const = 0; // logical cap
+  // Truncate to new_len (agent backtracking); false = cannot grow, or the
+  // target is older than an evicting layer still retains.
+  virtual bool rewind(int new_len) = 0;
+  virtual void clear() = 0; // reset for reuse
+};
+
+// A contiguous span of physical rows in a layer's pool.
+struct Run {
+  int start;
+  int len;
+};
+
+// 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.
+struct SeqStepPlan {
+  Run write[2];
+  int n_write;
+  Run read[2];
+  int n_read;
+  int read_base_pos;
+};
+
+// Backend face: per-layer layout for a step. `layer` selects the layer's
+// policy. nullopt = the step would exceed capacity or `layer` is out of range.
+class SequencePlanner {
+ public:
+  virtual ~SequencePlanner() = default;
+  virtual std::optional<SeqStepPlan> plan(int layer, int position, int T) = 0;
+};
+
+// Per-layer layout behavior (e.g. full-history flat). 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.
+  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
+  Kind kind = Kind::Flat;
+  int window =
+      0; // sliding-window size for windowing policies; 0 = full history
+};
+
+// Per-layer architecture facts + cache policy.
+struct LayerConfig {
+  LayerPolicy policy; // default Flat
+  int n_kv_heads;
+  int head_dim;
 };
 
-// Model facts a cache factory builds from. capacity is the logical cap;
-// n_layers is the number of attention layers; initial_capacity is the byte
-// layer's starting pool size before it grows (by doubling) toward capacity.
+// 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.
 struct CacheConfig {
   int capacity;
   int n_layers;
+  std::vector<LayerConfig> layers;
   int initial_capacity = 512;
 };
 

extension/llm/cache/cache_et.h

@@ -0,0 +1,64 @@
+/*
+ * 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
+
+// ExecuTorch adapter for the neutral cache core. The core (cache.h /
+// sequence_cache.h) is ET-independent and reports failures as
+// bool/std::optional so it is usable outside an ET runner. These thin inline
+// adapters map those results to ExecuTorch Error/Result (logging on failure)
+// for ET consumers -- the runner and the delegate byte layer. (The registry is
+// delegate-specific and already returns Result directly, so it needs no
+// adapter.)
+
+#include <optional>
+
+#include <executorch/extension/llm/cache/cache.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace cache {
+namespace et {
+
+using ::executorch::runtime::Error;
+using ::executorch::runtime::Result;
+
+// Plan a layer's step, or OutOfResources if it would exceed capacity (or the
+// layer is out of range).
+inline Result<SeqStepPlan>
+plan(SequencePlanner& planner, int layer, int position, int T) {
+  std::optional<SeqStepPlan> p = planner.plan(layer, position, T);
+  ET_CHECK_OR_RETURN_ERROR(
+      p.has_value(),
+      OutOfResources,
+      "cache: plan(layer=%d, position=%d, T=%d) exceeds capacity or bad layer",
+      layer,
+      position,
+      T);
+  return *p;
+}
+
+// Truncate the history, or InvalidArgument if new_len would grow it (or is
+// older than an evicting layer retains).
+inline Error rewind(SequenceControl& control, int new_len) {
+  ET_CHECK_OR_RETURN_ERROR(
+      control.rewind(new_len),
+      InvalidArgument,
+      "rewind: cannot grow to %d",
+      new_len);
+  return Error::Ok;
+}
+
+} // namespace et
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/cache/cache_registry.cpp

@@ -50,7 +50,7 @@ void CacheBuilderRegistry::register_builder(
     const std::string& kind,
     CacheBuilder builder) {
   std::lock_guard<std::mutex> lock(mu_);
-  builders_[backend_id + ":" + kind] = std::move(builder);
+  builders_[{backend_id, kind}] = std::move(builder);
 }
 
 Result<std::shared_ptr<CacheBase>> CacheBuilderRegistry::build(
@@ -60,7 +60,7 @@ Result<std::shared_ptr<CacheBase>> CacheBuilderRegistry::build(
   CacheBuilder builder;
   {
     std::lock_guard<std::mutex> lock(mu_);
-    const auto it = builders_.find(backend_id + ":" + kind);
+    const auto it = builders_.find({backend_id, kind});
     ET_CHECK_OR_RETURN_ERROR(
         it != builders_.end(),
         NotFound,

extension/llm/cache/cache_registry.h

@@ -12,15 +12,17 @@
 // is opaque to the host, so the runner (which knows the cache kind) creates the
 // cache and binds it to the delegate through a process-global registry; the two
 // sides rendezvous on a cache_key passed as a runtime backend-load option.
-// Caches are owned as CacheBase*, keeping the registry agnostic to the concrete
-// cache type. This layer is delegate-specific and may use ExecuTorch
-// Error/Result directly; the cache handle (cache.h) stays ET-free.
+// Caches are owned as CacheBase* and the faces are recovered via
+// as_control()/as_planner() (no RTTI). This layer is delegate-specific and may
+// use ExecuTorch Error/Result directly; the cache core (cache.h) stays ET-free.
 
 #include <functional>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
+#include <utility>
 
 #include <executorch/extension/llm/cache/cache.h>
 #include <executorch/runtime/core/error.h>
@@ -76,8 +78,8 @@ class CacheBuilderRegistry {
   CacheBuilderRegistry() = default;
 
   mutable std::mutex mu_;
-  std::unordered_map<std::string, CacheBuilder>
-      builders_; // backend_id + ":" + kind
+  std::map<std::pair<std::string, std::string>, CacheBuilder>
+      builders_; // keyed by (backend_id, kind)
 };
 
 // Process-global atomic counter -> "cache-N"; centralizes key generation so
@@ -86,7 +88,7 @@ std::string make_unique_key();
 
 // RAII: installs the cache into the global registry under a unique key on
 // construction and erases it on destruction (no leak on any exit path). Holds
-// the runner's shared_ptr for the lifetime of the generation loop.
+// the runner's shared_ptr and exposes the control face for the generation loop.
 class CacheSession {
  public:
   CacheSession(std::string key, std::shared_ptr<CacheBase> cache)
@@ -100,8 +102,8 @@ class CacheSession {
   CacheSession(const CacheSession&) = delete;
   CacheSession& operator=(const CacheSession&) = delete;
 
-  const std::shared_ptr<CacheBase>& cache() const {
-    return cache_;
+  SequenceControl* control() const {
+    return cache_->as_control();
   }
   const std::string& key() const {
     return key_;

extension/llm/cache/sequence_cache.h

@@ -0,0 +1,139 @@
+/*
+ * 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
+
+// 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.
+
+#include <algorithm>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include <executorch/extension/llm/cache/cache.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace cache {
+
+// Full history [0, length): one contiguous write run, read over all history.
+class FlatPolicy final : public LayoutPolicy {
+ public:
+  int retained_from(int /*length*/) const override {
+    return 0; // keeps all history
+  }
+  SeqStepPlan plan(int position, int T) const override {
+    const int end = position + T;
+    SeqStepPlan p{};
+    p.n_write = 1;
+    p.write[0] = Run{position, T}; // contiguous append, never wraps
+    p.n_read = 1;
+    p.read[0] = Run{0, end}; // attend over all history
+    p.read_base_pos = 0;
+    return p;
+  }
+};
+
+// 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.
+class SequenceCache : public CacheBase,
+                      public SequenceControl,
+                      public SequencePlanner {
+ public:
+  explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity) {
+    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.
+      const LayerConfig& lc =
+          cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l];
+      layer_to_policy_.push_back(policy_index(lc.policy));
+    }
+  }
+
+  // CacheBase: face recovery without RTTI.
+  SequenceControl* as_control() override {
+    return this;
+  }
+  SequencePlanner* as_planner() override {
+    return this;
+  }
+
+  // SequenceControl.
+  bool can_extend(int n = 1) const override {
+    return length_ + n <=
+        capacity_; // evicting layers reuse rows; capacity bounds
+  }
+  int capacity() const override {
+    return capacity_;
+  }
+  void clear() override {
+    length_ = 0;
+  }
+  bool rewind(int new_len) override {
+    if (new_len > length_) {
+      return false; // cannot grow
+    }
+    // An evicting layer physically drops everything older than it retains, so
+    // the target must be no older than the most-restrictive layer retains.
+    int floor = 0;
+    for (const auto& p : policies_) {
+      floor = std::max(floor, p->retained_from(length_));
+    }
+    if (new_len < floor) {
+      return false; // history evicted from an evicting layer
+    }
+    length_ = new_len;
+    return true;
+  }
+
+  // SequencePlanner.
+  std::optional<SeqStepPlan> plan(int layer, int position, int T) override {
+    if (layer < 0 || layer >= static_cast<int>(layer_to_policy_.size())) {
+      return std::nullopt;
+    }
+    const int end = position + T;
+    if (end > capacity_) {
+      return std::nullopt;
+    }
+    // Idempotent: plan runs once per layer per step, so commit the max, not a
+    // per-layer add.
+    length_ = std::max(length_, end);
+    return policies_[layer_to_policy_[layer]]->plan(position, T);
+  }
+
+ private:
+  int policy_index(const LayerPolicy& lp) {
+    for (std::size_t i = 0; i < specs_.size(); ++i) {
+      if (specs_[i].kind == lp.kind && specs_[i].window == lp.window) {
+        return static_cast<int>(i);
+      }
+    }
+    specs_.push_back(lp);
+    policies_.push_back(make_policy(lp));
+    return static_cast<int>(policies_.size() - 1);
+  }
+  std::unique_ptr<LayoutPolicy> make_policy(const LayerPolicy&) const {
+    return std::make_unique<FlatPolicy>();
+  }
+
+  int capacity_;
+  int length_ = 0;
+  std::vector<LayerPolicy> specs_; // parallel to policies_, for dedup
+  std::vector<std::unique_ptr<LayoutPolicy>> policies_;
+  std::vector<int> layer_to_policy_;
+};
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch