From 40c166664c40cdf707b2a56c4d249176a12cf29f Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 24 Jul 2026 07:11:25 -0700 Subject: [PATCH] Add off-graph KV-cache registry --- CMakeLists.txt | 2 + extension/llm/cache/CMakeLists.txt | 40 ++++++++ extension/llm/cache/cache.h | 40 ++++++++ extension/llm/cache/cache_registry.cpp | 83 +++++++++++++++++ extension/llm/cache/cache_registry.h | 118 ++++++++++++++++++++++++ extension/llm/cache/test/CMakeLists.txt | 17 ++++ extension/llm/cache/test/cache_test.cpp | 79 ++++++++++++++++ 7 files changed, 379 insertions(+) create mode 100644 extension/llm/cache/CMakeLists.txt create mode 100644 extension/llm/cache/cache.h create mode 100644 extension/llm/cache/cache_registry.cpp create mode 100644 extension/llm/cache/cache_registry.h create mode 100644 extension/llm/cache/test/CMakeLists.txt create mode 100644 extension/llm/cache/test/cache_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index b8d9b38f694..ff3b9e86f7e 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/extension/llm/cache/CMakeLists.txt b/extension/llm/cache/CMakeLists.txt new file mode 100644 index 00000000000..727ce68abc7 --- /dev/null +++ b/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() diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h new file mode 100644 index 00000000000..f60531dd62f --- /dev/null +++ b/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 diff --git a/extension/llm/cache/cache_registry.cpp b/extension/llm/cache/cache_registry.cpp new file mode 100644 index 00000000000..5ed463afc43 --- /dev/null +++ b/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 + +#include + +#include + +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 cache) { + std::lock_guard lock(mu_); + caches_[key] = std::move(cache); +} + +std::shared_ptr CacheRegistry::get(const std::string& key) const { + std::lock_guard 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 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 lock(mu_); + builders_[backend_id + ":" + kind] = std::move(builder); +} + +Result> CacheBuilderRegistry::build( + const std::string& backend_id, + const std::string& kind, + const CacheConfig& cfg) const { + CacheBuilder builder; + { + std::lock_guard 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 counter{0}; + return "cache-" + std::to_string(counter.fetch_add(1)); +} + +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/cache_registry.h b/extension/llm/cache/cache_registry.h new file mode 100644 index 00000000000..179b9e523ba --- /dev/null +++ b/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 +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { + +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +// Process-global map>. 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 cache); + std::shared_ptr get(const std::string& key) const; + void erase(const std::string& key); + + private: + CacheRegistry() = default; + + mutable std::mutex mu_; + std::unordered_map> 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(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> build( + const std::string& backend_id, + const std::string& kind, + const CacheConfig& cfg) const; + + private: + CacheBuilderRegistry() = default; + + mutable std::mutex mu_; + std::unordered_map + 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 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& cache() const { + return cache_; + } + const std::string& key() const { + return key_; + } + + private: + std::string key_; + std::shared_ptr cache_; +}; + +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/test/CMakeLists.txt b/extension/llm/cache/test/CMakeLists.txt new file mode 100644 index 00000000000..5161fb638fb --- /dev/null +++ b/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 +) diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp new file mode 100644 index 00000000000..1d238c23ed7 --- /dev/null +++ b/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 +#include + +#include + +#include +#include +#include + +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 cache = std::make_shared(); + 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(std::make_shared()); + }); + + 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()); + EXPECT_EQ(CacheRegistry::global().get(key), session.cache()); + } + EXPECT_EQ(CacheRegistry::global().get(key), nullptr); +}