Add the off-graph KV-cache cell layout to the neutral C++ layer

extension/llm/cache/CMakeLists.txt

@@ -14,7 +14,7 @@ if(NOT EXECUTORCH_ROOT)
   set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
 endif()
 
-add_library(extension_llm_cache cache_registry.cpp)
+add_library(extension_llm_cache cache_registry.cpp cell_cache.cpp)
 target_link_libraries(extension_llm_cache executorch_core)
 target_include_directories(
   extension_llm_cache PUBLIC ${_common_include_directories}

extension/llm/cache/cache.h

@@ -9,16 +9,12 @@
 #pragma once
 
 // 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 (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.
+// cache exposes a runner-facing control face and a backend-facing planner face,
+// recovered from the owning CacheBase*. Which pair it implements depends on the
+// layout: one sequence over per-layer runs, or many sequences over a pool of
+// per-token cells.
 
+#include <cstdint>
 #include <optional>
 #include <vector>
 
@@ -29,40 +25,54 @@ namespace cache {
 
 class SequenceControl;
 class SequencePlanner;
+class BatchControl;
+class CellStepper;
 
-// 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`).
+// Registry ownership anchor. A cache returns `this` from the faces it
+// implements and leaves the rest null.
 class CacheBase {
  public:
   virtual ~CacheBase() = default;
-  virtual SequenceControl* as_control() = 0;
-  virtual SequencePlanner* as_planner() = 0;
+  virtual SequenceControl* as_control() {
+    return nullptr;
+  }
+  virtual SequencePlanner* as_planner() {
+    return nullptr;
+  }
+  virtual BatchControl* as_batch_control() {
+    return nullptr;
+  }
+  virtual CellStepper* as_cell_stepper() {
+    return nullptr;
+  }
 };
 
-// Application (runner) face: lifecycle + admission, tensor-free.
-class SequenceControl {
+// Lifecycle and admission, tensor-free.
+class CacheControl {
  public:
-  virtual ~SequenceControl() = default;
+  virtual ~CacheControl() = 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
 };
 
+// Application face of a single-sequence cache: one length to rewind.
+class SequenceControl : public CacheControl {
+ public:
+  // Truncate to new_len; false = cannot grow, or the target is older than an
+  // evicting layer still retains.
+  virtual bool rewind(int new_len) = 0;
+};
+
 // 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 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.
+// Integer-only handoff to the backend byte layer. Runs are in logical order
+// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps.
+// read_base_pos is the logical position of read[0].start.
 struct SeqStepPlan {
   Run write[2];
   int n_write;
@@ -71,33 +81,61 @@ struct SeqStepPlan {
   int read_base_pos;
 };
 
-// Backend face. plan() is pure -- it computes a layer's layout for a step
-// without changing state; commit() advances the shared logical length once the
-// step is accepted. `layer` selects the layer's policy. nullopt = the step
-// would exceed capacity or `layer` is out of range.
+// Backend face. plan() is const: it computes a layer's layout without changing
+// state, and commit() advances the shared logical length. nullopt = the step
+// exceeds capacity, or `layer` is out of range.
 class SequencePlanner {
  public:
   virtual ~SequencePlanner() = default;
   virtual std::optional<SeqStepPlan> plan(int layer, int position, int T)
       const = 0;
-  // Advance the logical length past this step. Idempotent (commits the max), so
-  // calling it once per step -- not per layer -- suffices.
+  // Advance the logical length past this step. Idempotent, so once per step
+  // suffices.
   virtual void commit(const SeqStepPlan& plan) = 0;
 };
 
-// Per-layer layout behavior (flat = full history; ring = sliding window). Pure:
-// plan() has no side effects, so the controller (SequenceCache) owns length.
+// Per-layer layout: flat keeps all history, ring slides a window. Stateless.
 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).
+  // policy's window.
   virtual SeqStepPlan plan(int position, int T) const = 0;
-  // Oldest logical position still retained given the current length (0 for
-  // flat; length - window for ring). Used to bound rewind.
+  // Oldest logical position still retained at this length: 0 for flat,
+  // length - window for ring.
   virtual int retained_from(int length) const = 0;
 };
 
+// Application face of any multi-sequence cache: the sequence verbs. They run
+// between forwards, never during one.
+class BatchControl : public CacheControl {
+ public:
+  // Which sequence each of the next forward's tokens belongs to, one entry per
+  // token; every id must be one seq_new handed out. Also the admission gate:
+  // false = rejected and nothing changed, and a step that passes has room for
+  // its tokens. Whether its positions are well-formed is checked when the step
+  // is placed.
+  virtual bool declare_step(const std::vector<int32_t>& seq_ids) = 0;
+  // An id no live sequence is using, held until that sequence's last slot is
+  // freed. nullopt = every id is in use. Ids may also be chosen by the caller;
+  // this only guarantees the one it returns is not already taken.
+  virtual std::optional<int32_t> seq_new() = 0;
+  // A new sequence claiming src's slots below `upto`, all of them when unset.
+  // A shared slot keeps one position, so only a prefix can be shared, and the
+  // fork is a snapshot: slots src gains afterwards are its own. Nothing is
+  // copied. nullopt = an unknown or empty src, or no free sequence id.
+  virtual std::optional<int32_t> seq_clone(
+      int32_t src,
+      std::optional<int> upto) = 0;
+  // Drop the sequence's claim on positions [p0, p1). A slot frees only once
+  // no sequence owns it. False = an unknown sequence; a range owning nothing
+  // is a no-op.
+  virtual bool seq_rm(int32_t seq_id, int p0, std::optional<int> p1) = 0;
+  virtual int seq_len(int32_t seq_id) const = 0; // slots the sequence owns
+  // one past its newest position
+  virtual int next_pos(int32_t seq_id) const = 0;
+};
+
 // Per-layer cache kind and its parameters.
 struct LayerPolicy {
   enum class Kind : int {
@@ -115,29 +153,23 @@ struct LayerConfig {
   int head_dim;
 };
 
-// Model facts + runtime policy the byte layer sizes its pools from. capacity is
-// the logical cap; kv_dtype is the ET ScalarType the byte layer stores K/V in;
-// 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.
+// Model facts and the policy the byte layer sizes its pools from. `layers` is
+// per-layer: size 1 applies to every layer, else one entry each.
 struct CacheConfig {
-  int capacity;
+  int capacity; // logical cap in cells
   int n_layers;
   std::vector<LayerConfig> layers;
-  int kv_dtype;
-  int initial_capacity = 512;
+  int kv_dtype; // ET ScalarType the byte layer stores K/V in
+  int initial_capacity = 512; // starting pool size; grows lazily to capacity
+  // Max tokens per step; a ring layer sizes slots to window + max_write - 1.
+  // Unset = each ring layer uses its own window.
   std::optional<int> max_write;
 };
 
-// Whether `cfg` satisfies the contract above. Callers must check this before
-// constructing a cache: the `layers` broadcast rule is indexed directly, so a
-// list that is neither size 1 nor n_layers reads past the end. Reported as a
-// bool rather than thrown, so each backend picks its own failure mode.
+// Whether `cfg` satisfies the contract above.
 inline bool valid(const CacheConfig& cfg) {
-  // initial_capacity may be 0 (allocate nothing up front) but not negative, and
-  // may exceed capacity -- the byte layer clamps it.
+  // initial_capacity may be 0 but not negative, and may exceed capacity -- the
+  // byte layer clamps it.
   return cfg.capacity > 0 && cfg.n_layers > 0 && cfg.initial_capacity >= 0 &&
       (cfg.layers.size() == 1 ||
        cfg.layers.size() == static_cast<size_t>(cfg.n_layers));

extension/llm/cache/cache_registry.cpp

@@ -69,6 +69,14 @@ Result<std::shared_ptr<CacheBase>> CacheBuilderRegistry::build(
         kind.c_str());
     builder = it->second;
   }
+  // Checked here rather than in each cache: `layers` is indexed directly, so a
+  // list that is neither size 1 nor n_layers reads past the end.
+  ET_CHECK_OR_RETURN_ERROR(
+      valid(cfg),
+      InvalidArgument,
+      "cache: invalid CacheConfig for %s:%s",
+      backend_id.c_str(),
+      kind.c_str());
   return builder(cfg);
 }
 

extension/llm/cache/cache_registry.h

@@ -12,9 +12,8 @@
 // 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* 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.
+// Caches are owned as CacheBase* and the faces are recovered through its as_*
+// accessors (no RTTI), each null for a face the cache does not implement.
 
 #include <functional>
 #include <map>