Add off-graph KV-cache registry

CMakeLists.txt

@@ -903,6 +903,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
     )
   endif()
   list(APPEND _executorch_extensions tokenizers)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache)
+  list(APPEND _executorch_extensions extension_llm_cache)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)

extension/llm/cache/CMakeLists.txt

@@ -0,0 +1,40 @@
+# 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.
+
+# Install / rendezvous registry for the off-graph KV cache, shared across
+# backends. The neutral cache handle (cache.h) is ET-independent, so it is
+# usable outside an ET runner; the delegate-specific registry
+# (cache_registry.cpp) uses ExecuTorch Error/Result, hence the executorch_core
+# link.
+
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
+endif()
+
+add_library(extension_llm_cache cache_registry.cpp)
+target_link_libraries(extension_llm_cache executorch_core)
+target_include_directories(
+  extension_llm_cache PUBLIC ${_common_include_directories}
+)
+target_compile_options(extension_llm_cache PUBLIC ${_common_compile_options})
+
+install(
+  TARGETS extension_llm_cache
+  EXPORT ExecuTorchTargets
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  INCLUDES
+  DESTINATION ${_common_include_directories}
+)
+install(
+  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/cache
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+
+if(BUILD_TESTING)
+  add_subdirectory(test)
+endif()

extension/llm/cache/cache.h

@@ -0,0 +1,40 @@
+/*
+ * 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
+
+// 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.
+
+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 CacheBase {
+ public:
+  virtual ~CacheBase() = default;
+};
+
+// 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.
+struct CacheConfig {
+  int capacity;
+  int n_layers;
+  int initial_capacity = 512;
+};
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/cache/cache_registry.cpp

@@ -0,0 +1,83 @@
+/*
+ * 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/extension/llm/cache/cache_registry.h>
+
+#include <atomic>
+
+#include <executorch/runtime/core/error.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace cache {
+
+CacheRegistry& CacheRegistry::global() {
+  static CacheRegistry registry;
+  return registry;
+}
+
+void CacheRegistry::install(
+    const std::string& key,
+    std::shared_ptr<CacheBase> cache) {
+  std::lock_guard<std::mutex> lock(mu_);
+  caches_[key] = std::move(cache);
+}
+
+std::shared_ptr<CacheBase> CacheRegistry::get(const std::string& key) const {
+  std::lock_guard<std::mutex> lock(mu_);
+  const auto it = caches_.find(key);
+  return it == caches_.end() ? nullptr : it->second;
+}
+
+void CacheRegistry::erase(const std::string& key) {
+  std::lock_guard<std::mutex> lock(mu_);
+  caches_.erase(key);
+}
+
+CacheBuilderRegistry& CacheBuilderRegistry::global() {
+  static CacheBuilderRegistry registry;
+  return registry;
+}
+
+void CacheBuilderRegistry::register_builder(
+    const std::string& backend_id,
+    const std::string& kind,
+    CacheBuilder builder) {
+  std::lock_guard<std::mutex> lock(mu_);
+  builders_[backend_id + ":" + kind] = std::move(builder);
+}
+
+Result<std::shared_ptr<CacheBase>> CacheBuilderRegistry::build(
+    const std::string& backend_id,
+    const std::string& kind,
+    const CacheConfig& cfg) const {
+  CacheBuilder builder;
+  {
+    std::lock_guard<std::mutex> lock(mu_);
+    const auto it = builders_.find(backend_id + ":" + kind);
+    ET_CHECK_OR_RETURN_ERROR(
+        it != builders_.end(),
+        NotFound,
+        "no cache builder registered for %s:%s",
+        backend_id.c_str(),
+        kind.c_str());
+    builder = it->second;
+  }
+  return builder(cfg);
+}
+
+std::string make_unique_key() {
+  static std::atomic<uint64_t> counter{0};
+  return "cache-" + std::to_string(counter.fetch_add(1));
+}
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/cache/cache_registry.h

@@ -0,0 +1,118 @@
+/*
+ * 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
+
+// Install / rendezvous machinery for the off-graph KV cache. The DelegateHandle
+// 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.
+
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+
+#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 {
+
+using ::executorch::runtime::Error;
+using ::executorch::runtime::Result;
+
+// Process-global map<cache_key, shared_ptr<CacheBase>>. Ownership is shared:
+// the registry entry, the runner's session guard, and the delegate handle all
+// hold the cache, so erasing the entry mid-method is safe.
+class CacheRegistry {
+ public:
+  static CacheRegistry& global();
+
+  void install(const std::string& key, std::shared_ptr<CacheBase> cache);
+  std::shared_ptr<CacheBase> get(const std::string& key) const;
+  void erase(const std::string& key);
+
+ private:
+  CacheRegistry() = default;
+
+  mutable std::mutex mu_;
+  std::unordered_map<std::string, std::shared_ptr<CacheBase>> caches_;
+};
+
+// Cache kind is expressed by which factory you call: backends register a
+// builder per (backend_id, kind) and the kind survives only as an internal
+// lookup tag.
+using CacheBuilder =
+    std::function<std::shared_ptr<CacheBase>(const CacheConfig&)>;
+
+class CacheBuilderRegistry {
+ public:
+  static CacheBuilderRegistry& global();
+
+  void register_builder(
+      const std::string& backend_id,
+      const std::string& kind,
+      CacheBuilder builder);
+  // Returns Error::NotFound if no builder is registered for (backend_id, kind).
+  Result<std::shared_ptr<CacheBase>> build(
+      const std::string& backend_id,
+      const std::string& kind,
+      const CacheConfig& cfg) const;
+
+ private:
+  CacheBuilderRegistry() = default;
+
+  mutable std::mutex mu_;
+  std::unordered_map<std::string, CacheBuilder>
+      builders_; // backend_id + ":" + kind
+};
+
+// Process-global atomic counter -> "cache-N"; centralizes key generation so
+// keys never collide.
+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.
+class CacheSession {
+ public:
+  CacheSession(std::string key, std::shared_ptr<CacheBase> cache)
+      : key_(std::move(key)), cache_(std::move(cache)) {
+    CacheRegistry::global().install(key_, cache_);
+  }
+  ~CacheSession() {
+    CacheRegistry::global().erase(key_);
+  }
+
+  CacheSession(const CacheSession&) = delete;
+  CacheSession& operator=(const CacheSession&) = delete;
+
+  const std::shared_ptr<CacheBase>& cache() const {
+    return cache_;
+  }
+  const std::string& key() const {
+    return key_;
+  }
+
+ private:
+  std::string key_;
+  std::shared_ptr<CacheBase> cache_;
+};
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/cache/test/CMakeLists.txt

@@ -0,0 +1,17 @@
+# 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.
+
+cmake_minimum_required(VERSION 3.19)
+
+set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..)
+
+include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
+
+set(_test_srcs cache_test.cpp)
+
+et_cxx_test(
+  extension_llm_cache_test SOURCES ${_test_srcs} EXTRA_LIBS extension_llm_cache
+)

extension/llm/cache/test/cache_test.cpp

@@ -0,0 +1,79 @@
+/*
+ * 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/extension/llm/cache/cache.h>
+#include <executorch/extension/llm/cache/cache_registry.h>
+
+#include <memory>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <gtest/gtest.h>
+
+using executorch::extension::llm::cache::CacheBase;
+using executorch::extension::llm::cache::CacheBuilderRegistry;
+using executorch::extension::llm::cache::CacheConfig;
+using executorch::extension::llm::cache::CacheRegistry;
+using executorch::extension::llm::cache::CacheSession;
+using executorch::extension::llm::cache::make_unique_key;
+using executorch::runtime::Error;
+
+namespace {
+// Opaque CacheBase used to exercise the registry without a concrete cache
+// implementation (SequenceCache and its faces arrive in a later change).
+class StubCache : public CacheBase {};
+} // namespace
+
+// Initializes the ExecuTorch PAL so registry error paths (which ET_LOG) can
+// run.
+class CacheTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    executorch::runtime::runtime_init();
+  }
+};
+
+TEST_F(CacheTest, RegistryInstallGetErase) {
+  auto& reg = CacheRegistry::global();
+  const std::string key = make_unique_key();
+  EXPECT_EQ(reg.get(key), nullptr);
+
+  std::shared_ptr<CacheBase> cache = std::make_shared<StubCache>();
+  reg.install(key, cache);
+  EXPECT_EQ(reg.get(key), cache);
+
+  reg.erase(key);
+  EXPECT_EQ(reg.get(key), nullptr);
+}
+
+TEST_F(CacheTest, UniqueKeysDoNotCollide) {
+  EXPECT_NE(make_unique_key(), make_unique_key());
+}
+
+TEST_F(CacheTest, BuilderRegistryBuildsRegisteredKindElseError) {
+  auto& reg = CacheBuilderRegistry::global();
+  reg.register_builder("TestBackend", "stub", [](const CacheConfig&) {
+    return std::static_pointer_cast<CacheBase>(std::make_shared<StubCache>());
+  });
+
+  CacheConfig cfg{32, 1};
+  auto cache = reg.build("TestBackend", "stub", cfg);
+  ASSERT_TRUE(cache.ok());
+  EXPECT_NE(cache.get(), nullptr);
+
+  EXPECT_EQ(reg.build("TestBackend", "missing", cfg).error(), Error::NotFound);
+}
+
+TEST_F(CacheTest, SessionInstallsOnCtorErasesOnDtor) {
+  const std::string key = make_unique_key();
+  {
+    CacheSession session(key, std::make_shared<StubCache>());
+    EXPECT_EQ(CacheRegistry::global().get(key), session.cache());
+  }
+  EXPECT_EQ(CacheRegistry::global().get(key), nullptr);
+}