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..3f029724309 --- /dev/null +++ b/extension/llm/cache/CMakeLists.txt @@ -0,0 +1,38 @@ +# 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. + +# Neutral, tensor-free KV-cache core shared across backends (registry, +# bookkeeping, config). No backend or tensor dependency. CMake-only for now; +# buck targets are added later (see the buckify follow-up). + +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..169d1bf78ba --- /dev/null +++ b/extension/llm/cache/cache.h @@ -0,0 +1,155 @@ +/* + * 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, tensor-free KV-cache core shared across backends. A cache exposes +// two faces: a runner-facing control face (SequenceControl) and a +// backend-facing planner face (SequencePlanner). Both are recovered from the +// owning CacheBase* via as_control()/as_planner() -- static upcasts, so no +// dynamic_cast/RTTI and no diamond inheritance. Errors propagate as +// Error/Result (no exceptions, no abort). See extension/llm/cache for the +// full design. + +#include + +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { + +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +class SequenceControl; +class SequencePlanner; + +// Registry ownership / erasure anchor. The registry owns a cache as a +// CacheBase*; the runner recovers the control face and the backend recovers the +// planner face through these accessors (each concrete cache returns `this`). +class CacheBase { + public: + virtual ~CacheBase() = default; + virtual SequenceControl* as_control() = 0; + virtual SequencePlanner* as_planner() = 0; +}; + +// Application (runner) face: lifecycle + admission, tensor-free. +class SequenceControl { + public: + virtual ~SequenceControl() = default; + virtual bool can_extend(int n = 1) const = 0; // admission / hard-stop + virtual int capacity() const = 0; // logical cap + virtual Error rewind(int new_len) = 0; // truncate (agent backtracking) + virtual void clear() = 0; // reset for reuse +}; + +// Integer-only handoff from the planner to the backend byte layer: where to +// write this step and how much history to read. +struct SeqStepPlan { + int valid_len; // contiguous read bound + int write_start; // contiguous write run start + int write_len; // contiguous write run length +}; + +// Backend face: turns (position, T) into a write/read plan over the pool. +class SequencePlanner { + public: + virtual ~SequencePlanner() = default; + virtual Result plan(int position, int T) = 0; +}; + +// Model facts + runtime policy the byte layer sizes its pools from. +// Architecture facts (n_layers/n_kv_heads/head_dim) come from .pte metadata; +// capacity and initial_capacity are runtime policy. KV storage dtype joins this +// with the quantized pool; kept tensor-free here so the neutral core carries no +// backend types. +// +// n_kv_heads/head_dim are per-layer so hybrid / mixed-head / MLA models (KV +// dims vary by layer) need no config change; a single-element vector means +// uniform across all layers (the common case). +struct CacheConfig { + int capacity; + int n_layers; + std::vector n_kv_heads; + std::vector head_dim; + // Initial pool size in cells: the byte layer's lazy-doubling pool starts at + // initial_capacity and doubles up to `capacity`, so short sequences avoid + // reallocation. + int initial_capacity = 512; +}; + +// Flat single-sequence cache: non-wrapping history [0, length) tracked by a +// length counter. Implements both faces; a backend byte layer builds on it. +class FlatCache : public CacheBase, + public SequenceControl, + public SequencePlanner { + public: + explicit FlatCache(int capacity) : capacity_(capacity) {} + + // CacheBase: face recovery without RTTI. + SequenceControl* as_control() override { + return this; + } + SequencePlanner* as_planner() override { + return this; + } + + // SequenceControl. + bool can_extend(int n = 1) const override { + return length_ + n <= capacity_; + } + int capacity() const override { + return capacity_; + } + // Rows are overwritten on the next write, so reuse/truncation is byte-free. + void clear() override { + length_ = 0; + } + Error rewind(int new_len) override { + ET_CHECK_OR_RETURN_ERROR( + new_len <= length_, + InvalidArgument, + "rewind: cannot grow (new_len=%d > length=%d)", + new_len, + length_); + length_ = new_len; + return Error::Ok; + } + + // SequencePlanner: place T tokens at `position`, return the write run + read + // bound. + Result plan(int position, int T) override { + const int end = position + T; + ET_CHECK_OR_RETURN_ERROR( + end <= capacity_, + InvalidArgument, + "cache: %d cells exceeds capacity %d", + end, + capacity_); + if (end > length_) { + length_ = end; + } + return SeqStepPlan{ + /*valid_len=*/end, + /*write_start=*/position, + /*write_len=*/T}; + } + + private: + int capacity_; + int length_ = 0; +}; + +} // 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..f38bdd7cbe9 --- /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(), + InvalidArgument, + "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..d4948946913 --- /dev/null +++ b/extension/llm/cache/cache_registry.h @@ -0,0 +1,113 @@ +/* + * 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. +// Neutral: caches are owned as CacheBase*, and the faces are recovered via +// as_control()/as_planner() (no RTTI). + +#include +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { + +// 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::InvalidArgument 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 and exposes the control face for 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; + + SequenceControl* control() const { + return cache_->as_control(); + } + 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..6d2acd6aa3d --- /dev/null +++ b/extension/llm/cache/test/cache_test.cpp @@ -0,0 +1,125 @@ +/* + * 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::FlatCache; +using executorch::extension::llm::cache::make_unique_key; +using executorch::runtime::Error; + +// Initializes the ExecuTorch PAL so error paths (which ET_LOG) can run. +class CacheTest : public ::testing::Test { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + } +}; + +TEST_F(CacheTest, PlanAppendsAndAdvancesLength) { + FlatCache cache(/*capacity=*/8); + auto p0 = cache.plan(/*position=*/0, /*T=*/4); // prefill + ASSERT_TRUE(p0.ok()); + EXPECT_EQ(p0->write_start, 0); + EXPECT_EQ(p0->write_len, 4); + EXPECT_EQ(p0->valid_len, 4); + + auto p1 = cache.plan(/*position=*/4, /*T=*/1); // decode + ASSERT_TRUE(p1.ok()); + EXPECT_EQ(p1->write_start, 4); + EXPECT_EQ(p1->valid_len, 5); +} + +TEST_F(CacheTest, CanExtendBoundsAtCapacity) { + FlatCache cache(/*capacity=*/2); + EXPECT_TRUE(cache.can_extend(2)); + ASSERT_TRUE(cache.plan(/*position=*/0, /*T=*/2).ok()); + EXPECT_FALSE(cache.can_extend(1)); + EXPECT_EQ( + cache.plan(/*position=*/2, /*T=*/1).error(), Error::InvalidArgument); +} + +TEST_F(CacheTest, RewindTruncatesAndClearResets) { + FlatCache cache(/*capacity=*/8); + ASSERT_TRUE(cache.plan(0, 5).ok()); + EXPECT_EQ(cache.rewind(2), Error::Ok); + EXPECT_TRUE(cache.can_extend(6)); + EXPECT_EQ(cache.rewind(9), Error::InvalidArgument); + cache.clear(); + EXPECT_TRUE(cache.can_extend(8)); +} + +TEST_F(CacheTest, FaceRecoveryReturnsSameObject) { + FlatCache cache(/*capacity=*/4); + CacheBase* base = &cache; + ASSERT_NE(base->as_control(), nullptr); + ASSERT_NE(base->as_planner(), nullptr); + EXPECT_TRUE(base->as_control()->can_extend(4)); + auto plan = base->as_planner()->plan(0, 1); + ASSERT_TRUE(plan.ok()); + EXPECT_EQ(plan->valid_len, 1); +} + +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(16); + reg.install(key, cache); + EXPECT_EQ(reg.get(key), cache); + EXPECT_TRUE(reg.get(key)->as_control()->can_extend(16)); + + reg.erase(key); + EXPECT_EQ(reg.get(key), nullptr); +} + +TEST_F(CacheTest, UniqueKeysDoNotCollide) { + EXPECT_NE(make_unique_key(), make_unique_key()); +} + +TEST_F(CacheTest, BuilderRegistryBuildsRegisteredKindElseErrors) { + auto& reg = CacheBuilderRegistry::global(); + reg.register_builder("TestBackend", "flat", [](const CacheConfig& cfg) { + return std::static_pointer_cast( + std::make_shared(cfg.capacity)); + }); + + CacheConfig cfg{ + /*capacity=*/32, + /*n_layers=*/1, + /*n_kv_heads=*/{1}, + /*head_dim=*/{8}}; + auto cache = reg.build("TestBackend", "flat", cfg); + ASSERT_TRUE(cache.ok()); + EXPECT_EQ(cache.get()->as_control()->capacity(), 32); + + EXPECT_EQ( + reg.build("TestBackend", "missing", cfg).error(), Error::InvalidArgument); +} + +TEST_F(CacheTest, SessionInstallsOnCtorErasesOnDtor) { + const std::string key = make_unique_key(); + { + CacheSession session(key, std::make_shared(4)); + EXPECT_NE(CacheRegistry::global().get(key), nullptr); + EXPECT_TRUE(session.control()->can_extend(4)); + } + EXPECT_EQ(CacheRegistry::global().get(key), nullptr); +}