From a7438ce58300c32c7eb56516c29b06e5eff6c4db Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Mon, 27 Jul 2026 15:02:26 -0700 Subject: [PATCH] Add off-graph KV-cache SequenceCache and flat policy --- extension/llm/cache/cache.h | 101 +++++++++++++++-- extension/llm/cache/cache_et.h | 64 +++++++++++ extension/llm/cache/cache_registry.cpp | 4 +- extension/llm/cache/cache_registry.h | 18 +-- extension/llm/cache/sequence_cache.h | 139 ++++++++++++++++++++++++ extension/llm/cache/test/cache_test.cpp | 108 +++++++++++++++--- 6 files changed, 401 insertions(+), 33 deletions(-) create mode 100644 extension/llm/cache/cache_et.h create mode 100644 extension/llm/cache/sequence_cache.h diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index f60531dd62f..a79cebe3675 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -8,29 +8,112 @@ #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. +// Neutral, tensor-free, ET-independent KV-cache core shared across backends. A +// cache exposes two faces recovered from the owning CacheBase* via +// as_control()/as_planner() (static upcasts -- no dynamic_cast/RTTI, no +// diamond): a runner-facing control face (SequenceControl) and a backend-facing +// planner face (SequencePlanner). One controller (SequenceCache) drives all +// layers and dispatches per-layer layout to a LayoutPolicy (full-history flat +// today), so a multi-layer model shares one logical length. Errors are plain +// C++ (bool / std::optional); cache_et.h adapts them to Error/Result for ET +// consumers. + +#include +#include 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 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 + // Truncate to new_len (agent backtracking); false = cannot grow, or the + // target is older than an evicting layer still retains. + virtual bool rewind(int new_len) = 0; + virtual void clear() = 0; // reset for reuse +}; + +// A contiguous span of physical rows in a layer's pool. +struct Run { + int start; + int len; +}; + +// Integer-only handoff from the planner to the backend byte layer. Runs are in +// logical order (oldest -> newest); a flat layer uses one run, a windowing +// layer up to two (a write/read that wraps its buffer splits in two). +// read_base_pos is the logical position of read[0].start (0 for flat), so the +// backend can align RoPE / the attention mask. +struct SeqStepPlan { + Run write[2]; + int n_write; + Run read[2]; + int n_read; + int read_base_pos; +}; + +// Backend face: per-layer layout for a step. `layer` selects the layer's +// policy. nullopt = the step would exceed capacity or `layer` is out of range. +class SequencePlanner { + public: + virtual ~SequencePlanner() = default; + virtual std::optional plan(int layer, int position, int T) = 0; +}; + +// Per-layer layout behavior (e.g. full-history flat). Pure: plan() has no side +// effects, so the controller (SequenceCache) owns length. +class LayoutPolicy { + public: + virtual ~LayoutPolicy() = default; + // Write/read runs for T cells at logical `position`. Precondition: T fits the + // policy's window (the runner chunks prefill so a step fits). + virtual SeqStepPlan plan(int position, int T) const = 0; + // Oldest logical position still retained given the current length (0 for a + // full-history policy; a windowing policy retains only its last window). Used + // to bound rewind. + virtual int retained_from(int length) const = 0; +}; + +// Per-layer cache kind and its parameters. +struct LayerPolicy { + enum class Kind : int { Flat = 0 }; // serialized values: append-only + Kind kind = Kind::Flat; + int window = + 0; // sliding-window size for windowing policies; 0 = full history +}; + +// Per-layer architecture facts + cache policy. +struct LayerConfig { + LayerPolicy policy; // default Flat + int n_kv_heads; + int head_dim; }; -// 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. +// Model facts + runtime policy the byte layer sizes its pools from. capacity is +// the logical cap; initial_capacity tunes the byte layer's lazy-doubling pool. +// `layers` is per-layer: size 1 == uniform across all layers, else == n_layers. struct CacheConfig { int capacity; int n_layers; + std::vector layers; int initial_capacity = 512; }; diff --git a/extension/llm/cache/cache_et.h b/extension/llm/cache/cache_et.h new file mode 100644 index 00000000000..dabf9cc916a --- /dev/null +++ b/extension/llm/cache/cache_et.h @@ -0,0 +1,64 @@ +/* + * 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 + +// ExecuTorch adapter for the neutral cache core. The core (cache.h / +// sequence_cache.h) is ET-independent and reports failures as +// bool/std::optional so it is usable outside an ET runner. These thin inline +// adapters map those results to ExecuTorch Error/Result (logging on failure) +// for ET consumers -- the runner and the delegate byte layer. (The registry is +// delegate-specific and already returns Result directly, so it needs no +// adapter.) + +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { +namespace et { + +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +// Plan a layer's step, or OutOfResources if it would exceed capacity (or the +// layer is out of range). +inline Result +plan(SequencePlanner& planner, int layer, int position, int T) { + std::optional p = planner.plan(layer, position, T); + ET_CHECK_OR_RETURN_ERROR( + p.has_value(), + OutOfResources, + "cache: plan(layer=%d, position=%d, T=%d) exceeds capacity or bad layer", + layer, + position, + T); + return *p; +} + +// Truncate the history, or InvalidArgument if new_len would grow it (or is +// older than an evicting layer retains). +inline Error rewind(SequenceControl& control, int new_len) { + ET_CHECK_OR_RETURN_ERROR( + control.rewind(new_len), + InvalidArgument, + "rewind: cannot grow to %d", + new_len); + return Error::Ok; +} + +} // namespace et +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/cache_registry.cpp b/extension/llm/cache/cache_registry.cpp index 5ed463afc43..225868e68cf 100644 --- a/extension/llm/cache/cache_registry.cpp +++ b/extension/llm/cache/cache_registry.cpp @@ -50,7 +50,7 @@ void CacheBuilderRegistry::register_builder( const std::string& kind, CacheBuilder builder) { std::lock_guard lock(mu_); - builders_[backend_id + ":" + kind] = std::move(builder); + builders_[{backend_id, kind}] = std::move(builder); } Result> CacheBuilderRegistry::build( @@ -60,7 +60,7 @@ Result> CacheBuilderRegistry::build( CacheBuilder builder; { std::lock_guard lock(mu_); - const auto it = builders_.find(backend_id + ":" + kind); + const auto it = builders_.find({backend_id, kind}); ET_CHECK_OR_RETURN_ERROR( it != builders_.end(), NotFound, diff --git a/extension/llm/cache/cache_registry.h b/extension/llm/cache/cache_registry.h index 179b9e523ba..1fbc438756d 100644 --- a/extension/llm/cache/cache_registry.h +++ b/extension/llm/cache/cache_registry.h @@ -12,15 +12,17 @@ // 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. +// Caches are owned as CacheBase* and the faces are recovered via +// as_control()/as_planner() (no RTTI). This layer is delegate-specific and may +// use ExecuTorch Error/Result directly; the cache core (cache.h) stays ET-free. #include +#include #include #include #include #include +#include #include #include @@ -76,8 +78,8 @@ class CacheBuilderRegistry { CacheBuilderRegistry() = default; mutable std::mutex mu_; - std::unordered_map - builders_; // backend_id + ":" + kind + std::map, CacheBuilder> + builders_; // keyed by (backend_id, kind) }; // Process-global atomic counter -> "cache-N"; centralizes key generation so @@ -86,7 +88,7 @@ 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. +// 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) @@ -100,8 +102,8 @@ class CacheSession { CacheSession(const CacheSession&) = delete; CacheSession& operator=(const CacheSession&) = delete; - const std::shared_ptr& cache() const { - return cache_; + SequenceControl* control() const { + return cache_->as_control(); } const std::string& key() const { return key_; diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h new file mode 100644 index 00000000000..5c9ac170bc3 --- /dev/null +++ b/extension/llm/cache/sequence_cache.h @@ -0,0 +1,139 @@ +/* + * 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 + +// The neutral single-sequence controller (SequenceCache) and its flat layout +// policy (FlatPolicy). SequenceCache owns the one logical length for the whole +// model and dispatches per-layer layout to a policy, so a multi-layer model +// stays coherent. Tensor-free / ET-independent. + +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { + +// Full history [0, length): one contiguous write run, read over all history. +class FlatPolicy final : public LayoutPolicy { + public: + int retained_from(int /*length*/) const override { + return 0; // keeps all history + } + SeqStepPlan plan(int position, int T) const override { + const int end = position + T; + SeqStepPlan p{}; + p.n_write = 1; + p.write[0] = Run{position, T}; // contiguous append, never wraps + p.n_read = 1; + p.read[0] = Run{0, end}; // attend over all history + p.read_base_pos = 0; + return p; + } +}; + +// One controller for all layers: owns the single logical length, admission, and +// rewind; dispatches per-layer layout to a shared LayoutPolicy. Policies are +// deduped by (kind, window), so a uniform model holds a single policy object. +class SequenceCache : public CacheBase, + public SequenceControl, + public SequencePlanner { + public: + explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity) { + layer_to_policy_.reserve(cfg.n_layers); + for (int l = 0; l < cfg.n_layers; ++l) { + // layers size 1 = one config broadcast to every layer, else per-layer. + const LayerConfig& lc = + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; + layer_to_policy_.push_back(policy_index(lc.policy)); + } + } + + // 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_; // evicting layers reuse rows; capacity bounds + } + int capacity() const override { + return capacity_; + } + void clear() override { + length_ = 0; + } + bool rewind(int new_len) override { + if (new_len > length_) { + return false; // cannot grow + } + // An evicting layer physically drops everything older than it retains, so + // the target must be no older than the most-restrictive layer retains. + int floor = 0; + for (const auto& p : policies_) { + floor = std::max(floor, p->retained_from(length_)); + } + if (new_len < floor) { + return false; // history evicted from an evicting layer + } + length_ = new_len; + return true; + } + + // SequencePlanner. + std::optional plan(int layer, int position, int T) override { + if (layer < 0 || layer >= static_cast(layer_to_policy_.size())) { + return std::nullopt; + } + const int end = position + T; + if (end > capacity_) { + return std::nullopt; + } + // Idempotent: plan runs once per layer per step, so commit the max, not a + // per-layer add. + length_ = std::max(length_, end); + return policies_[layer_to_policy_[layer]]->plan(position, T); + } + + private: + int policy_index(const LayerPolicy& lp) { + for (std::size_t i = 0; i < specs_.size(); ++i) { + if (specs_[i].kind == lp.kind && specs_[i].window == lp.window) { + return static_cast(i); + } + } + specs_.push_back(lp); + policies_.push_back(make_policy(lp)); + return static_cast(policies_.size() - 1); + } + std::unique_ptr make_policy(const LayerPolicy&) const { + return std::make_unique(); + } + + int capacity_; + int length_ = 0; + std::vector specs_; // parallel to policies_, for dedup + std::vector> policies_; + std::vector layer_to_policy_; +}; + +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index 1d238c23ed7..6db56cc5333 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -7,7 +7,9 @@ */ #include +#include #include +#include #include @@ -20,17 +22,21 @@ 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::LayerConfig; +using executorch::extension::llm::cache::LayerPolicy; using executorch::extension::llm::cache::make_unique_key; +using executorch::extension::llm::cache::SequenceCache; using executorch::runtime::Error; +namespace et = executorch::extension::llm::cache::et; 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 {}; +LayerConfig flat_layer() { + return LayerConfig{LayerPolicy{LayerPolicy::Kind::Flat, 0}, 2, 8}; +} } // namespace -// Initializes the ExecuTorch PAL so registry error paths (which ET_LOG) can -// run. +// Initializes the ExecuTorch PAL so the ET adapter's error paths (which ET_LOG) +// can run. class CacheTest : public ::testing::Test { protected: void SetUp() override { @@ -38,14 +44,69 @@ class CacheTest : public ::testing::Test { } }; +// ---- Flat policy ----------------------------------------------------------- + +TEST_F(CacheTest, FlatPlanAppendsAndReadsAllHistory) { + SequenceCache cache(CacheConfig{8, 1, {flat_layer()}}); + + auto p0 = cache.plan(/*layer=*/0, /*position=*/0, /*T=*/4); // prefill + ASSERT_TRUE(p0.has_value()); + EXPECT_EQ(p0->n_write, 1); + EXPECT_EQ(p0->write[0].start, 0); + EXPECT_EQ(p0->write[0].len, 4); + EXPECT_EQ(p0->n_read, 1); + EXPECT_EQ(p0->read[0].start, 0); + EXPECT_EQ(p0->read[0].len, 4); + EXPECT_EQ(p0->read_base_pos, 0); + + auto p1 = cache.plan(0, 4, 1); // decode + ASSERT_TRUE(p1.has_value()); + EXPECT_EQ(p1->write[0].start, 4); + EXPECT_EQ(p1->read[0].len, 5); +} + +// ---- Admission / rewind ---------------------------------------------------- + +TEST_F(CacheTest, CanExtendBoundedByCapacity) { + SequenceCache cache(CacheConfig{2, 1, {flat_layer()}}); + EXPECT_TRUE(cache.can_extend(2)); + ASSERT_TRUE(cache.plan(0, 0, 2).has_value()); + EXPECT_FALSE(cache.can_extend(1)); + EXPECT_FALSE(cache.plan(0, 2, 1).has_value()); // exceeds capacity +} + +TEST_F(CacheTest, FlatRewindsFreelyToZero) { + SequenceCache cache(CacheConfig{8, 1, {flat_layer()}}); + ASSERT_TRUE(cache.plan(0, 0, 5).has_value()); + EXPECT_TRUE(cache.rewind(0)); // flat retains all history + EXPECT_FALSE(cache.rewind(1)); // cannot grow + cache.clear(); + EXPECT_TRUE(cache.can_extend(8)); +} + +// ---- Faces / registry / session -------------------------------------------- + +TEST_F(CacheTest, FaceRecoveryReturnsSameObject) { + SequenceCache cache(CacheConfig{4, 1, {flat_layer()}}); + 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, 0, 1); + ASSERT_TRUE(plan.has_value()); + EXPECT_EQ(plan->read[0].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(); + std::shared_ptr cache = + std::make_shared(CacheConfig{16, 1, {flat_layer()}}); 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); @@ -55,16 +116,17 @@ TEST_F(CacheTest, UniqueKeysDoNotCollide) { EXPECT_NE(make_unique_key(), make_unique_key()); } -TEST_F(CacheTest, BuilderRegistryBuildsRegisteredKindElseError) { +TEST_F(CacheTest, BuilderBuildsRegisteredKindElseError) { auto& reg = CacheBuilderRegistry::global(); - reg.register_builder("TestBackend", "stub", [](const CacheConfig&) { - return std::static_pointer_cast(std::make_shared()); + reg.register_builder("TestBackend", "seq", [](const CacheConfig& cfg) { + return std::static_pointer_cast( + std::make_shared(cfg)); }); - CacheConfig cfg{32, 1}; - auto cache = reg.build("TestBackend", "stub", cfg); + CacheConfig cfg{32, 1, {flat_layer()}}; + auto cache = reg.build("TestBackend", "seq", cfg); ASSERT_TRUE(cache.ok()); - EXPECT_NE(cache.get(), nullptr); + EXPECT_EQ(cache.get()->as_control()->capacity(), 32); EXPECT_EQ(reg.build("TestBackend", "missing", cfg).error(), Error::NotFound); } @@ -72,8 +134,26 @@ TEST_F(CacheTest, BuilderRegistryBuildsRegisteredKindElseError) { 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()); + CacheSession session( + key, + std::make_shared(CacheConfig{4, 1, {flat_layer()}})); + EXPECT_NE(CacheRegistry::global().get(key), nullptr); + EXPECT_TRUE(session.control()->can_extend(4)); } EXPECT_EQ(CacheRegistry::global().get(key), nullptr); } + +// ---- ET adapter (maps core bool/optional to Error/Result) ------------------ + +TEST_F(CacheTest, EtAdapterMapsResultsAndCodes) { + SequenceCache cache(CacheConfig{2, 1, {flat_layer()}}); + auto ok = et::plan(cache, /*layer=*/0, /*position=*/0, /*T=*/2); + ASSERT_TRUE(ok.ok()); + EXPECT_EQ(ok->read[0].len, 2); + EXPECT_EQ( + et::plan(cache, 0, 2, 1).error(), Error::OutOfResources); // over capacity + EXPECT_FALSE(et::plan(cache, 5, 0, 1).ok()); // bad layer + + EXPECT_EQ(et::rewind(cache, 9), Error::InvalidArgument); // cannot grow + EXPECT_EQ(et::rewind(cache, 1), Error::Ok); +}