[Native] Serve kvcache::update_and_attend in the ggml engine

CMakeLists.txt

@@ -1045,6 +1045,12 @@ if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)
   list(APPEND _executorch_extensions extension_runner_util)
 endif()
 
+# The native ggml engine serves off-graph KV caches from the cache registry.
+if(EXECUTORCH_BUILD_NATIVE_GGML AND NOT TARGET extension_llm_cache)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache)
+  list(APPEND _executorch_extensions extension_llm_cache)
+endif()
+
 if(EXECUTORCH_BUILD_EXTENSION_TENSOR)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/tensor)
   list(APPEND _executorch_extensions extension_tensor)

backends/native/runtime/engine/ggml/CMakeLists.txt

@@ -32,6 +32,7 @@ target_link_libraries(executorch_ggml INTERFACE ggml)
 
 add_library(
   native_ggml_engine
+  GgmlCache.cpp
   GgmlEngine.cpp
   GgmlExecutor.cpp
   GgmlGraphBuilder.cpp
@@ -47,7 +48,7 @@ add_library(executorch::native_ggml_engine ALIAS native_ggml_engine)
 target_link_libraries(
   native_ggml_engine
   PUBLIC native_engine
-  PRIVATE executorch_ggml
+  PRIVATE executorch_ggml extension_llm_cache
 )
 target_compile_options(
   native_ggml_engine PRIVATE $<$<CXX_COMPILER_ID:MSVC>:/EHsc>

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

@@ -0,0 +1,181 @@
+// 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.
+
+#include <executorch/backends/native/runtime/engine/ggml/GgmlCache.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <exception>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include <executorch/backends/native/runtime/engine/ggml/GgmlEngine.h>
+#include <executorch/backends/native/runtime/graph/ScalarType.h>
+#include <executorch/extension/llm/cache/cache_registry.h>
+#include <executorch/runtime/platform/log.h>
+
+namespace ptn {
+
+namespace cache = executorch::extension::llm::cache;
+
+namespace {
+
+ggml_type kv_type(int kv_dtype) {
+  switch (static_cast<ScalarType>(kv_dtype)) {
+    case ScalarType::Float:
+      return GGML_TYPE_F32;
+    case ScalarType::Half:
+      return GGML_TYPE_F16;
+    default:
+      throw std::invalid_argument(
+          "ggml cache: kv_dtype must be Float or Half, got " +
+          std::to_string(kv_dtype));
+  }
+}
+
+const cache::CacheGeometry& checked(
+    const cache::CacheGeometry& geometry,
+    const cache::CacheConfig& cfg) {
+  if (!cache::valid(geometry, cfg)) {
+    throw std::invalid_argument("ggml cache: invalid geometry or config");
+  }
+  if (geometry.dflash_block_size) {
+    throw std::invalid_argument("ggml cache: DFlash geometry is unsupported");
+  }
+  for (const cache::LayerGeometry& layer : geometry.layers) {
+    if (layer.policy.kind != cache::LayerPolicy::Kind::Flat) {
+      throw std::invalid_argument(
+          "ggml cache: only flat (full-history) layers are supported");
+    }
+  }
+  return geometry;
+}
+
+} // namespace
+
+GgmlSequenceCache::GgmlSequenceCache(
+    const cache::CacheGeometry& geometry,
+    const cache::CacheConfig& cfg)
+    : SequenceCache(checked(geometry, cfg), cfg),
+      geometry_(geometry),
+      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) {
+  ggml_backend_buffer_type_t buffer_type =
+      ggml_backend_get_default_buffer_type(backend);
+  if (storage_ && buffer_type != buffer_type_) {
+    throw std::runtime_error(
+        std::string("ggml cache: storage lives in ") +
+        ggml_backend_buft_name(buffer_type_) + ", not " +
+        ggml_backend_buft_name(buffer_type));
+  }
+  if (rows > capacity()) {
+    throw std::out_of_range(
+        "ggml cache: " + std::to_string(rows) +
+        " rows exceed the capacity of " + std::to_string(capacity()));
+  }
+  if (storage_ && rows <= storage_->rows) {
+    return;
+  }
+  int64_t target = storage_ ? storage_->rows : initial_rows_;
+  while (target < rows) {
+    target *= 2;
+  }
+  auto next = std::make_shared<Storage>();
+  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),
+      .mem_buffer = nullptr,
+      .no_alloc = true,
+  };
+  next->context.reset(ggml_init(params));
+  if (!next->context) {
+    throw std::runtime_error("ggml cache: failed to create tensor context");
+  }
+  for (size_t i = 0; i < geometry_.layers.size(); ++i) {
+    const cache::LayerGeometry& layer = geometry_.layers[i];
+    Layer out;
+    for (ggml_tensor** t : {&out.k, &out.v}) {
+      *t = ggml_new_tensor_3d(
+          next->context.get(),
+          type_,
+          layer.head_dim,
+          next->rows,
+          layer.n_kv_heads);
+      ggml_format_name(*t, "kv_cache.%s.%zu", t == &out.k ? "k" : "v", i);
+    }
+    next->layers.push_back(out);
+  }
+  next->buffer.reset(
+      ggml_backend_alloc_ctx_tensors_from_buft(next->context.get(), buffer_type));
+  if (!next->buffer) {
+    throw std::runtime_error(
+        "ggml cache: failed to allocate K/V storage of " +
+        std::to_string(next->rows) + " rows");
+  }
+  ggml_backend_buffer_clear(next->buffer.get(), 0);
+
+  // Each head's rows are contiguous, but the head stride grows with the rows.
+  std::vector<uint8_t> staging;
+  for (size_t i = 0; storage_ && i < storage_->layers.size(); ++i) {
+    const Layer& old = storage_->layers[i];
+    for (const auto& [from, to] :
+         {std::pair{old.k, next->layers[i].k},
+          std::pair{old.v, next->layers[i].v}}) {
+      staging.resize(ggml_nbytes(from));
+      ggml_backend_tensor_get(from, staging.data(), 0, staging.size());
+      const size_t head_bytes = static_cast<size_t>(from->ne[1]) * from->nb[1];
+      for (int64_t h = 0; h < from->ne[2]; ++h) {
+        ggml_backend_tensor_set(
+            to,
+            staging.data() + h * from->nb[2],
+            h * to->nb[2],
+            head_bytes);
+      }
+    }
+  }
+
+  buffer_type_ = buffer_type;
+  storage_ = std::move(next);
+}
+
+void* GgmlSequenceCache::face(cache::FaceId id) {
+  return cache::expose<
+      cache::SequenceControl,
+      cache::SequencePlanner,
+      GgmlSequenceCache>(this, id);
+}
+
+namespace {
+
+// CacheFactory::build reports a null cache as an error, so a cache the engine
+// cannot serve comes back as one rather than as an exception.
+template <typename CacheT>
+cache::CacheBuilder builder() {
+  return [](const cache::CacheGeometry& geometry,
+            const cache::CacheConfig& cfg) -> std::shared_ptr<cache::Cache> {
+    try {
+      return std::make_shared<CacheT>(geometry, cfg);
+    } catch (const std::exception& error) {
+      ET_LOG(Error, "%s", error.what());
+      return nullptr;
+    }
+  };
+}
+
+} // 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>());
+}
+
+} // namespace ptn

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

@@ -0,0 +1,78 @@
+// 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
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <ggml-backend.h>
+#include <ggml-cpp.h>
+#include <ggml.h>
+
+#include <executorch/extension/llm/cache/sequence_cache.h>
+
+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]); row r
+// holds logical position r. Storage starts at the config's initial_capacity
+// rows and doubles toward capacity as steps need more. 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";
+
+  struct Layer {
+    ggml_tensor* k = nullptr;
+    ggml_tensor* v = nullptr;
+  };
+
+  // One allocation of K/V storage. A regrow moves the rows into a new one;
+  // the old one lives on while any executable still holds it.
+  struct Storage {
+    ggml_context_ptr context;
+    ggml_backend_buffer_ptr buffer;
+    std::vector<Layer> layers;
+    int rows = 0;
+  };
+
+  // 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;
+
+  // Ensures storage 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
+  // storage stay valid while their executable holds it.
+  const std::shared_ptr<const Storage>& storage() const {
+    return storage_;
+  }
+
+ protected:
+  void* face(executorch::extension::llm::cache::FaceId id) override;
+
+ private:
+  executorch::extension::llm::cache::CacheGeometry geometry_;
+  ggml_type type_;
+  int initial_rows_;
+  ggml_backend_buffer_type_t buffer_type_ = nullptr;
+  std::shared_ptr<const Storage> storage_;
+};
+
+} // namespace ptn

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

@@ -16,7 +16,9 @@
 #include <executorch/backends/native/runtime/Method.h>
 #include <executorch/backends/native/runtime/Program.h>
 #include <executorch/backends/native/runtime/deserialize/Package.h>
+#include <executorch/backends/native/runtime/engine/ggml/GgmlCache.h>
 #include <executorch/backends/native/runtime/engine/ggml/GgmlExecutor.h>
+#include <executorch/extension/llm/cache/cache_registry.h>
 
 namespace ptn {
 
@@ -81,17 +83,52 @@ std::unique_ptr<EngineContext> GgmlEngineHost::create_context(
   return create_ggml_context(std::move(program), std::move(package));
 }
 
+std::unique_ptr<EngineContext> GgmlEngineHost::create_context_with_options(
+    std::shared_ptr<const Program> program,
+    std::shared_ptr<const Package> package,
+    const EngineOptions& options) {
+  return create_ggml_context(std::move(program), std::move(package), options);
+}
+
 std::unique_ptr<GgmlEngineContext> GgmlEngineHost::create_ggml_context(
     std::shared_ptr<const Program> program,
-    std::shared_ptr<const Package> package) {
+    std::shared_ptr<const Package> package,
+    const EngineOptions& options) {
+  namespace cache = executorch::extension::llm::cache;
+  std::shared_ptr<GgmlSequenceCache> 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);
+    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) {
+      throw std::invalid_argument(
+          impl_->name + ": KV cache '" + value +
+          "' was not built for this engine (CacheFactory backend '" +
+          kGgmlCacheBackendId + "', kind 'single')");
+    }
+    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(program),
+      std::move(package),
+      impl_->options,
+      std::move(kv_cache)));
 }
 
 struct GgmlEngineContext::Impl {
-  explicit Impl(GgmlEngineOptions value) : options(std::move(value)) {}
+  Impl(GgmlEngineOptions value, std::shared_ptr<GgmlSequenceCache> cache)
+      : options(std::move(value)), kv_cache(std::move(cache)) {}
 
   GgmlEngineOptions options;
+  std::shared_ptr<GgmlSequenceCache> kv_cache;
   size_t materialized_constant_bytes = 0;
   size_t heap_copied_constant_bytes = 0;
   GgmlMutableKeyOwners mutable_key_owners;
@@ -100,9 +137,10 @@ struct GgmlEngineContext::Impl {
 GgmlEngineContext::GgmlEngineContext(
     std::shared_ptr<const Program> program,
     std::shared_ptr<const Package> package,
-    GgmlEngineOptions options)
+    GgmlEngineOptions options,
+    std::shared_ptr<GgmlSequenceCache> cache)
     : EngineContext(std::move(program), std::move(package)),
-      impl_(std::make_unique<Impl>(std::move(options))) {}
+      impl_(std::make_unique<Impl>(std::move(options), std::move(cache))) {}
 
 GgmlEngineContext::~GgmlEngineContext() = default;
 
@@ -114,7 +152,8 @@ std::unique_ptr<EngineExecutable> GgmlEngineContext::compile_method(
       impl_->options,
       &impl_->materialized_constant_bytes,
       &impl_->heap_copied_constant_bytes,
-      impl_->mutable_key_owners);
+      impl_->mutable_key_owners,
+      impl_->kv_cache.get());
 }
 
 size_t GgmlEngineContext::materialized_constant_bytes() const {