[Native] Separate the ggml engine's K/V storage from its sequence cache

backends/native/runtime/engine/ggml/GgmlCache.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -12,6 +12,7 @@
 #include <memory>
 #include <stdexcept>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <executorch/backends/native/runtime/engine/ggml/GgmlEngine.h>
@@ -70,16 +71,30 @@
 
 } // namespace
 
+GgmlKvPool::GgmlKvPool(
+    const cache::CacheGeometry& geometry,
+    int kv_dtype,
+    int initial_rows,
+    int capacity,
+    std::vector<int> ring_rows)
+    : geometry_(geometry),
+      ring_rows_(std::move(ring_rows)),
+      type_(kv_type(kv_dtype)),
+      initial_rows_(std::clamp(initial_rows, 1, capacity)),
+      capacity_(capacity) {}
+
 GgmlSequenceCache::GgmlSequenceCache(
     const cache::CacheGeometry& geometry,
     const cache::CacheConfig& cfg)
     : SequenceCache(checked(geometry, cfg), cfg),
-      geometry_(geometry),
-      ring_rows_(ring_rows(geometry, cfg)),
-      type_(kv_type(cfg.kv_dtype)),
-      initial_rows_(std::clamp(cfg.initial_capacity, 1, cfg.capacity)) {}
-
-void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) {
+      pool_(
+          geometry,
+          cfg.kv_dtype,
+          cfg.initial_capacity,
+          cfg.capacity,
+          ring_rows(geometry, cfg)) {}
+
+void GgmlKvPool::reserve(ggml_backend_t backend, int rows) {
   ggml_backend_buffer_type_t buffer_type =
       ggml_backend_get_default_buffer_type(backend);
   if (storage_ && buffer_type != buffer_type_) {
@@ -88,10 +103,10 @@
         ggml_backend_buft_name(buffer_type_) + ", not " +
         ggml_backend_buft_name(buffer_type));
   }
-  if (rows > capacity()) {
+  if (rows > capacity_) {
     throw std::out_of_range(
         "ggml cache: " + std::to_string(rows) +
-        " rows exceed the capacity of " + std::to_string(capacity()));
+        " rows exceed the capacity of " + std::to_string(capacity_));
   }
   if (storage_ && rows <= storage_->rows) {
     return;
@@ -101,7 +116,7 @@
     target *= 2;
   }
   auto next = std::make_shared<Storage>();
-  next->rows = static_cast<int>(std::min<int64_t>(target, capacity()));
+  next->rows = static_cast<int>(std::min<int64_t>(target, capacity_));
 
   ggml_init_params params{
       .mem_size = ggml_tensor_overhead() * (2 * geometry_.layers.size() + 1),
@@ -114,7 +129,7 @@
   }
   for (size_t i = 0; i < geometry_.layers.size(); ++i) {
     const cache::LayerGeometry& layer = geometry_.layers[i];
-    Layer out{.window = layer.policy.window};
+    Layer out{.window = layer.policy.window, .ring = ring_rows_[i] > 0};
     for (ggml_tensor** t : {&out.k, &out.v}) {
       *t = ggml_new_tensor_3d(
           next->context.get(),
@@ -186,7 +201,7 @@
 
 } // namespace
 
 void register_ggml_cache_builders() {
   // Registering twice is harmless; the first builder stays.
   (void)cache::CacheFactory::global().register_builder(
       kGgmlCacheBackendId, cache::kind::kSingle, builder<GgmlSequenceCache>());

backends/native/runtime/engine/ggml/GgmlCache.h

@@ -19,25 +19,20 @@
 
 namespace ptn {
 
-// The off-graph KV cache of one sequence for the ggml engine: the neutral
-// SequenceCache controller plus K/V storage per layer in the buffer type of
-// the first backend that binds it. Each layer's K and V are ggml
-// [head_dim, rows, n_kv_heads] (torch [1, n_kv_heads, rows, head_dim]). A flat
-// layer's row r holds logical position r; its storage starts at the config's
-// initial_capacity rows and doubles toward capacity as steps need more. A ring
-// layer holds its window + max_write - 1 rows from the start, and the
-// SequencePlanner says which rows a step writes and reads. Like SequenceCache,
-// it is not thread-safe: executables bound to one cache must be serialized.
-class GgmlSequenceCache final
-    : public executorch::extension::llm::cache::SequenceCache {
+// K/V storage per cache layer in the buffer type of the first backend that
+// binds it. Each layer's K and V are ggml [head_dim, rows, n_kv_heads] (torch
+// [1, n_kv_heads, rows, head_dim]). A growing layer starts at `initial_rows`
+// and doubles toward `capacity` as steps need more; a ring layer holds its
+// fixed rows from the start. Not thread-safe.
+class GgmlKvPool {
  public:
-  static constexpr const char* kFaceName = "ptn.ggml.SequenceCache";
-
   struct Layer {
     ggml_tensor* k = nullptr;
     ggml_tensor* v = nullptr;
-    // Sliding window of a ring layer; 0 for a flat one.
+    // The layer's sliding window, 0 for none; layers sharing one share a mask.
     int window = 0;
+    // A ring layer attends all its fixed rows; a growing one, a prefix.
+    bool ring = false;
   };
 
   // One allocation of K/V storage. A regrow moves the rows into a new one;
@@ -46,21 +41,23 @@ class GgmlSequenceCache final
     ggml_context_ptr context;
     ggml_backend_buffer_ptr buffer;
     std::vector<Layer> layers;
-    // Rows of every flat layer.
+    // Rows of every growing layer.
     int rows = 0;
   };
 
-  // Throws on geometry or config the ggml engine cannot serve.
-  GgmlSequenceCache(
+  // `ring_rows` holds each layer's fixed rows, or 0 for a growing layer.
+  // Throws for a kv_dtype other than Float or Half.
+  GgmlKvPool(
       const executorch::extension::llm::cache::CacheGeometry& geometry,
-      const executorch::extension::llm::cache::CacheConfig& cfg);
-
-  GgmlSequenceCache(const GgmlSequenceCache&) = delete;
-  GgmlSequenceCache& operator=(const GgmlSequenceCache&) = delete;
+      int kv_dtype,
+      int initial_rows,
+      int capacity,
+      std::vector<int> ring_rows);
 
-  // Ensures flat layers of at least `rows` rows in `backend`'s default buffer
-  // type, moving the rows it holds into a larger allocation if needed. Throws
-  // if `rows` exceeds capacity or an earlier call used another buffer type.
+  // Ensures growing layers of at least `rows` rows in `backend`'s default
+  // buffer type, moving the rows it holds into a larger allocation if needed.
+  // Throws if `rows` exceeds capacity or an earlier call used another buffer
+  // type.
   void reserve(ggml_backend_t backend, int rows);
 
   // The current storage; null before the first reserve. Graphs built over a
@@ -69,17 +66,44 @@ class GgmlSequenceCache final
     return storage_;
   }
 
- protected:
-  void* face(executorch::extension::llm::cache::FaceId id) override;
-
  private:
   executorch::extension::llm::cache::CacheGeometry geometry_;
-  // Per layer: the fixed rows of a ring layer, 0 for a flat one.
   std::vector<int> ring_rows_;
   ggml_type type_;
   int initial_rows_;
+  int capacity_;
   ggml_backend_buffer_type_t buffer_type_ = nullptr;
   std::shared_ptr<const Storage> storage_;
 };
 
+// The off-graph KV cache of one sequence for the ggml engine: the neutral
+// SequenceCache controller over a GgmlKvPool. A flat layer's row r holds
+// logical position r; a ring layer holds its window + max_write - 1 rows, and
+// the SequencePlanner says which rows a step writes and reads. Like
+// SequenceCache, it is not thread-safe: executables bound to one cache must be
+// serialized.
+class GgmlSequenceCache final
+    : public executorch::extension::llm::cache::SequenceCache {
+ public:
+  static constexpr const char* kFaceName = "ptn.ggml.SequenceCache";
+
+  // Throws on geometry or config the ggml engine cannot serve.
+  GgmlSequenceCache(
+      const executorch::extension::llm::cache::CacheGeometry& geometry,
+      const executorch::extension::llm::cache::CacheConfig& cfg);
+
+  GgmlSequenceCache(const GgmlSequenceCache&) = delete;
+  GgmlSequenceCache& operator=(const GgmlSequenceCache&) = delete;
+
+  GgmlKvPool& pool() {
+    return pool_;
+  }
+
+ protected:
+  void* face(executorch::extension::llm::cache::FaceId id) override;
+
+ private:
+  GgmlKvPool pool_;
+};
+
 } // namespace ptn

backends/native/runtime/engine/ggml/GgmlExecutor.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -1293,19 +1293,22 @@
   // the graphs built over the old storage, which this executable keeps alive
   // until then.
   void bind_offgraph(int64_t rows) {
-    kv_cache_->reserve(primary_backend_, static_cast<int>(rows));
-    if (kv_cache_->storage() == offgraph_storage_) {
+    kv_cache_->pool().reserve(primary_backend_, static_cast<int>(rows));
+    if (kv_cache_->pool().storage() == offgraph_storage_) {
       return;
     }
     for (CachedGraph& cached : cached_graphs_) {
       ggml_backend_sched_synchronize(cached.scheduler.get());
     }
     cached_graphs_.clear();
-    offgraph_storage_ = kv_cache_->storage();
+    offgraph_storage_ = kv_cache_->pool().storage();
     offgraph_cache_.layers.clear();
-    for (const GgmlSequenceCache::Layer& layer : offgraph_storage_->layers) {
+    for (const GgmlKvPool::Layer& layer : offgraph_storage_->layers) {
       offgraph_cache_.layers.push_back(
-          {.k = layer.k, .v = layer.v, .window = layer.window});
+          {.k = layer.k,
+           .v = layer.v,
+           .window = layer.window,
+           .ring = layer.ring});
     }
   }
 
@@ -1393,7 +1396,7 @@
       std::vector<int64_t> positions(static_cast<size_t>(group.keys), -1);
       int64_t position = plan.read_base_pos;
       for (int r = 0; r < plan.n_read; ++r) {
         const Run run = plan.read[r];
         if (run.start < 0 || run.start + run.len > group.keys) {
           throw std::logic_error(
               "ggml: KV cache read exceeds the graph's keys");
@@ -1509,7 +1512,7 @@
   OffGraphCacheBinding offgraph_cache_;
   // The storage offgraph_cache_ and cached graphs point into. Declared before
   // cached_graphs_ so it outlives them.
-  std::shared_ptr<const GgmlSequenceCache::Storage> offgraph_storage_;
+  std::shared_ptr<const GgmlKvPool::Storage> offgraph_storage_;
   // The step execute() is running, checked once.
   std::optional<OffGraphStep> offgraph_step_;
   std::vector<InputSlot> inputs_;

backends/native/runtime/engine/ggml/GgmlGraphBuilder.h

@@ -94,12 +94,13 @@ struct OffGraphCacheBinding {
     // K and V storage, ggml [head_dim, rows, n_kv_heads].
     ggml_tensor* k = nullptr;
     ggml_tensor* v = nullptr;
-    // Sliding window of a ring layer, which attends all its rows; 0 for a
-    // flat layer.
+    // The layer's sliding window, 0 for none; layers sharing one share a mask.
     int window = 0;
+    // A ring layer attends all its rows; any other, the first `keys`.
+    bool ring = false;
   };
   std::vector<Layer> layers;
-  // Rows a flat layer's attention reads, a prefix of its storage.
+  // Rows a non-ring layer's attention reads, a prefix of its storage.
   int64_t keys = 0;
 };
 

backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -398,8 +398,8 @@
 }
 
 // Writes this step's K/V into the cache rows the executor uploads, then
-// attends the layer's rows under the uploaded mask: a flat layer's prefix, or
-// all of a ring layer. Q/K/V are [B,H,T,D].
+// attends the layer's rows under the uploaded mask: a prefix, or all of a
+// ring layer. Q/K/V are [B,H,T,D].
 void GraphBuilder::lower_update_and_attend(const Node& node) {
   if (offgraph_cache_ == nullptr) {
     throw std::runtime_error(
@@ -417,7 +417,7 @@
         "ggml: update_and_attend layer_id " + std::to_string(layer) +
         " is outside the cache's " + std::to_string(num_layers) + " layers");
   }
-  const auto [k_cache, v_cache, window] =
+  const auto [k_cache, v_cache, window, ring] =
       offgraph_cache_->layers[static_cast<size_t>(layer)];
   if (query.sizes.size() != 4 || key.sizes.size() != 4 ||
       val.sizes.size() != 4 || query.sizes[0] != 1 || key.sizes[0] != 1 ||
@@ -437,7 +437,7 @@
         "ggml: update_and_attend layer " + std::to_string(layer) +
         " K/V shape disagrees with the cache geometry");
   }
-  const int64_t keys = window > 0 ? k_cache->ne[1] : offgraph_cache_->keys;
+  const int64_t keys = ring ? k_cache->ne[1] : offgraph_cache_->keys;
   ggml_context* ctx = instance_.context.get();
   if (!instance_.offgraph) {
     instance_.offgraph = OffGraphUploads{.tokens = tokens};