[Native] Serve many sequences from one ggml engine KV cache

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

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -182,6 +182,22 @@
       GgmlSequenceCache>(this, id);
 }
 
+GgmlCellCache::GgmlCellCache(
+    const cache::CacheGeometry& geometry,
+    const cache::CacheConfig& cfg)
+    : CellCache(checked(geometry, cfg), cfg),
+      pool_(
+          geometry,
+          cfg.kv_dtype,
+          cfg.initial_capacity,
+          cfg.capacity,
+          std::vector<int>(geometry.layers.size(), 0)) {}
+
+void* GgmlCellCache::face(cache::FaceId id) {
+  return cache::
+      expose<cache::BatchControl, cache::CellStepper, GgmlCellCache>(this, id);
+}
+
 namespace {
 
 // CacheFactory::build reports a null cache as an error, so a cache the engine
@@ -201,10 +217,14 @@
 
 } // 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>());
+  for (const char* kind : {cache::kind::kBatched, cache::kind::kBatchedCell}) {
+    (void)cache::CacheFactory::global().register_builder(
+        kGgmlCacheBackendId, kind, builder<GgmlCellCache>());
+  }
 }
 
 } // namespace ptn

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

@@ -15,6 +15,7 @@
 #include <ggml-cpp.h>
 #include <ggml.h>
 
+#include <executorch/extension/llm/cache/cell_cache.h>
 #include <executorch/extension/llm/cache/sequence_cache.h>
 
 namespace ptn {
@@ -106,4 +107,51 @@ class GgmlSequenceCache final
   GgmlKvPool pool_;
 };
 
+// The off-graph KV cache of many sequences for the ggml engine: the neutral
+// CellCache over a GgmlKvPool whose rows are its cells, every layer growing
+// toward the pool's capacity. A step's cells and masks come from place_step; a
+// windowed layer keeps every cell and its mask narrows to the window. Like
+// CellCache, it is not thread-safe: executables bound to one cache must be
+// serialized.
+class GgmlCellCache final
+    : public executorch::extension::llm::cache::CellCache {
+ public:
+  static constexpr const char* kFaceName = "ptn.ggml.CellCache";
+
+  // Throws on geometry or config the ggml engine cannot serve.
+  GgmlCellCache(
+      const executorch::extension::llm::cache::CacheGeometry& geometry,
+      const executorch::extension::llm::cache::CacheConfig& cfg);
+
+  GgmlCellCache(const GgmlCellCache&) = delete;
+  GgmlCellCache& operator=(const GgmlCellCache&) = delete;
+
+  GgmlKvPool& pool() {
+    return pool_;
+  }
+
+ protected:
+  void* face(executorch::extension::llm::cache::FaceId id) override;
+
+ private:
+  GgmlKvPool pool_;
+};
+
+// The off-graph cache an executable binds: its storage and its layout, one of
+// a single sequence or cells.
+struct GgmlKvCache {
+  GgmlSequenceCache* sequence = nullptr;
+  GgmlCellCache* cells = nullptr;
+
+  explicit operator bool() const {
+    return sequence != nullptr || cells != nullptr;
+  }
+  GgmlKvPool& pool() const {
+    return sequence != nullptr ? sequence->pool() : cells->pool();
+  }
+  int capacity() const {
+    return sequence != nullptr ? sequence->capacity() : cells->capacity();
+  }
+};
+
 } // namespace ptn

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

@@ -95,40 +95,47 @@ std::unique_ptr<GgmlEngineContext> GgmlEngineHost::create_ggml_context(
     std::shared_ptr<const Package> package,
     const EngineOptions& options) {
   namespace cache = executorch::extension::llm::cache;
-  std::shared_ptr<GgmlSequenceCache> kv_cache;
+  std::shared_ptr<cache::Cache> installed;
+  GgmlKvCache kv_cache;
   for (const auto& [name, value] : options) {
     if (name != cache::kCacheKeyOption) {
       throw std::invalid_argument(
           impl_->name + ": unsupported context option '" + name + "'");
     }
-    std::shared_ptr<cache::Cache> installed =
-        cache::CacheRegistry::global().get(value);
+    installed = cache::CacheRegistry::global().get(value);
     if (!installed) {
       throw std::invalid_argument(
           impl_->name + ": no KV cache is installed under key '" + value + "'");
     }
-    GgmlSequenceCache* ggml_cache = installed->as<GgmlSequenceCache>();
-    if (ggml_cache == nullptr) {
+    kv_cache = {
+        .sequence = installed->as<GgmlSequenceCache>(),
+        .cells = installed->as<GgmlCellCache>()};
+    if (!kv_cache) {
       throw std::invalid_argument(
           impl_->name + ": KV cache '" + value +
           "' was not built for this engine (CacheFactory backend '" +
-          kGgmlCacheBackendId + "', kind 'single')");
+          kGgmlCacheBackendId + "')");
     }
-    kv_cache = std::shared_ptr<GgmlSequenceCache>(installed, ggml_cache);
   }
   return std::unique_ptr<GgmlEngineContext>(new GgmlEngineContext(
       std::move(program),
       std::move(package),
       impl_->options,
-      std::move(kv_cache)));
+      std::move(installed),
+      kv_cache));
 }
 
 struct GgmlEngineContext::Impl {
-  Impl(GgmlEngineOptions value, std::shared_ptr<GgmlSequenceCache> cache)
-      : options(std::move(value)), kv_cache(std::move(cache)) {}
+  Impl(
+      GgmlEngineOptions value,
+      std::shared_ptr<executorch::extension::llm::cache::Cache> owner,
+      GgmlKvCache cache)
+      : options(std::move(value)), cache(std::move(owner)), kv_cache(cache) {}
 
   GgmlEngineOptions options;
-  std::shared_ptr<GgmlSequenceCache> kv_cache;
+  // Keeps the cache kv_cache points into alive.
+  std::shared_ptr<executorch::extension::llm::cache::Cache> cache;
+  GgmlKvCache kv_cache;
   size_t materialized_constant_bytes = 0;
   size_t heap_copied_constant_bytes = 0;
   GgmlMutableKeyOwners mutable_key_owners;
@@ -138,9 +145,13 @@ GgmlEngineContext::GgmlEngineContext(
     std::shared_ptr<const Program> program,
     std::shared_ptr<const Package> package,
     GgmlEngineOptions options,
-    std::shared_ptr<GgmlSequenceCache> cache)
+    std::shared_ptr<executorch::extension::llm::cache::Cache> cache,
+    GgmlKvCache kv_cache)
     : EngineContext(std::move(program), std::move(package)),
-      impl_(std::make_unique<Impl>(std::move(options), std::move(cache))) {}
+      impl_(std::make_unique<Impl>(
+          std::move(options),
+          std::move(cache),
+          kv_cache)) {}
 
 GgmlEngineContext::~GgmlEngineContext() = default;
 
@@ -153,7 +164,7 @@ std::unique_ptr<EngineExecutable> GgmlEngineContext::compile_method(
       &impl_->materialized_constant_bytes,
       &impl_->heap_copied_constant_bytes,
       impl_->mutable_key_owners,
-      impl_->kv_cache.get());
+      impl_->kv_cache);
 }
 
 size_t GgmlEngineContext::materialized_constant_bytes() const {

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

@@ -12,12 +12,17 @@
 
 #include <executorch/backends/native/runtime/engine/Engine.h>
 
+namespace executorch::extension::llm::cache {
+class Cache;
+} // namespace executorch::extension::llm::cache
+
 namespace ptn {
 
 // CacheFactory backend id of the ggml engine's off-graph KV caches. A method
 // that calls kvcache::update_and_attend compiles only in a context created
 // with the kCacheKeyOption ("llm_cache_registry_key") option naming a cache
-// built under this id with kind::kSingle.
+// built under this id: kind::kSingle for one sequence, or kind::kBatched /
+// kind::kBatchedCell for many sequences over shared cells.
 inline constexpr char kGgmlCacheBackendId[] = "NativeGgml";
 
 // Registers the ggml cache builders with CacheFactory::global(). Idempotent.
@@ -52,7 +57,7 @@ struct GgmlEngineOptions {
 };
 
 class GgmlEngineContext;
-class GgmlSequenceCache;
+struct GgmlKvCache;
 
 class GgmlEngineHost final : public EngineHost {
  public:
@@ -68,7 +73,7 @@ class GgmlEngineHost final : public EngineHost {
       std::shared_ptr<const Package> package) override;
 
   // Accepts kCacheKeyOption, resolved through CacheRegistry::global(). Contexts
-  // bound to one cache share its sequence: the caller must serialize them.
+  // bound to one cache share its sequences: the caller must serialize them.
   std::unique_ptr<EngineContext> create_context_with_options(
       std::shared_ptr<const Program> program,
       std::shared_ptr<const Package> package,
@@ -105,7 +110,8 @@ class GgmlEngineContext final : public EngineContext {
       std::shared_ptr<const Program> program,
       std::shared_ptr<const Package> package,
       GgmlEngineOptions options,
-      std::shared_ptr<GgmlSequenceCache> cache);
+      std::shared_ptr<executorch::extension::llm::cache::Cache> cache,
+      GgmlKvCache kv_cache);
 
   std::unique_ptr<EngineExecutable> compile_method(
       const Method& method) override;