diff --git a/extension/llm/cache/CMakeLists.txt b/extension/llm/cache/CMakeLists.txt index 727ce68abc7..8e3be66f5d7 100644 --- a/extension/llm/cache/CMakeLists.txt +++ b/extension/llm/cache/CMakeLists.txt @@ -14,7 +14,7 @@ if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) endif() -add_library(extension_llm_cache cache_registry.cpp) +add_library(extension_llm_cache cache_registry.cpp cell_cache.cpp) target_link_libraries(extension_llm_cache executorch_core) target_include_directories( extension_llm_cache PUBLIC ${_common_include_directories} diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 4194fe6a665..fb25e43f4a6 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -9,16 +9,12 @@ #pragma once // 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 (flat = full -// history, ring = sliding window), so a mixed model (e.g. gemma4's alternating -// full/sliding-window layers) shares one logical length. Errors are plain C++ -// (bool / std::optional); cache_et.h adapts them to Error/Result for ET -// consumers. +// cache exposes a runner-facing control face and a backend-facing planner face, +// recovered from the owning CacheBase*. Which pair it implements depends on the +// layout: one sequence over per-layer runs, or many sequences over a pool of +// per-token cells. +#include #include #include @@ -29,40 +25,54 @@ namespace cache { class SequenceControl; class SequencePlanner; +class BatchControl; +class CellStepper; -// 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`). +// Registry ownership anchor. A cache returns `this` from the faces it +// implements and leaves the rest null. class CacheBase { public: virtual ~CacheBase() = default; - virtual SequenceControl* as_control() = 0; - virtual SequencePlanner* as_planner() = 0; + virtual SequenceControl* as_control() { + return nullptr; + } + virtual SequencePlanner* as_planner() { + return nullptr; + } + virtual BatchControl* as_batch_control() { + return nullptr; + } + virtual CellStepper* as_cell_stepper() { + return nullptr; + } }; -// Application (runner) face: lifecycle + admission, tensor-free. -class SequenceControl { +// Lifecycle and admission, tensor-free. +class CacheControl { public: - virtual ~SequenceControl() = default; + virtual ~CacheControl() = 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 }; +// Application face of a single-sequence cache: one length to rewind. +class SequenceControl : public CacheControl { + public: + // Truncate to new_len; false = cannot grow, or the target is older than an + // evicting layer still retains. + virtual bool rewind(int new_len) = 0; +}; + // 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 ring layer up -// to two (a write/read that wraps the buffer splits in two). read_base_pos is -// the logical position of read[0].start (0 for flat; the window start for -// ring), so the backend can align RoPE / the attention mask. +// Integer-only handoff to the backend byte layer. Runs are in logical order +// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps. +// read_base_pos is the logical position of read[0].start. struct SeqStepPlan { Run write[2]; int n_write; @@ -71,33 +81,61 @@ struct SeqStepPlan { int read_base_pos; }; -// Backend face. plan() is pure -- it computes a layer's layout for a step -// without changing state; commit() advances the shared logical length once the -// step is accepted. `layer` selects the layer's policy. nullopt = the step -// would exceed capacity or `layer` is out of range. +// Backend face. plan() is const: it computes a layer's layout without changing +// state, and commit() advances the shared logical length. nullopt = the step +// exceeds capacity, or `layer` is out of range. class SequencePlanner { public: virtual ~SequencePlanner() = default; virtual std::optional plan(int layer, int position, int T) const = 0; - // Advance the logical length past this step. Idempotent (commits the max), so - // calling it once per step -- not per layer -- suffices. + // Advance the logical length past this step. Idempotent, so once per step + // suffices. virtual void commit(const SeqStepPlan& plan) = 0; }; -// Per-layer layout behavior (flat = full history; ring = sliding window). Pure: -// plan() has no side effects, so the controller (SequenceCache) owns length. +// Per-layer layout: flat keeps all history, ring slides a window. Stateless. 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). + // policy's window. virtual SeqStepPlan plan(int position, int T) const = 0; - // Oldest logical position still retained given the current length (0 for - // flat; length - window for ring). Used to bound rewind. + // Oldest logical position still retained at this length: 0 for flat, + // length - window for ring. virtual int retained_from(int length) const = 0; }; +// Application face of any multi-sequence cache: the sequence verbs. They run +// between forwards, never during one. +class BatchControl : public CacheControl { + public: + // Which sequence each of the next forward's tokens belongs to, one entry per + // token; every id must be one seq_new handed out. Also the admission gate: + // false = rejected and nothing changed, and a step that passes has room for + // its tokens. Whether its positions are well-formed is checked when the step + // is placed. + virtual bool declare_step(const std::vector& seq_ids) = 0; + // An id no live sequence is using, held until that sequence's last slot is + // freed. nullopt = every id is in use. Ids may also be chosen by the caller; + // this only guarantees the one it returns is not already taken. + virtual std::optional seq_new() = 0; + // A new sequence claiming src's slots below `upto`, all of them when unset. + // A shared slot keeps one position, so only a prefix can be shared, and the + // fork is a snapshot: slots src gains afterwards are its own. Nothing is + // copied. nullopt = an unknown or empty src, or no free sequence id. + virtual std::optional seq_clone( + int32_t src, + std::optional upto) = 0; + // Drop the sequence's claim on positions [p0, p1). A slot frees only once + // no sequence owns it. False = an unknown sequence; a range owning nothing + // is a no-op. + virtual bool seq_rm(int32_t seq_id, int p0, std::optional p1) = 0; + virtual int seq_len(int32_t seq_id) const = 0; // slots the sequence owns + // one past its newest position + virtual int next_pos(int32_t seq_id) const = 0; +}; + // Per-layer cache kind and its parameters. struct LayerPolicy { enum class Kind : int { @@ -115,29 +153,23 @@ struct LayerConfig { int head_dim; }; -// Model facts + runtime policy the byte layer sizes its pools from. capacity is -// the logical cap; kv_dtype is the ET ScalarType the byte layer stores K/V in; -// initial_capacity tunes the byte layer's lazy-doubling pool; max_write is the -// max tokens written per step (a ring layer sizes its slots to window + -// max_write - 1 so a multi-token step fits); unset means each ring layer uses -// its own window. `layers` is per-layer: size 1 == uniform across all layers, -// else == n_layers. +// Model facts and the policy the byte layer sizes its pools from. `layers` is +// per-layer: size 1 applies to every layer, else one entry each. struct CacheConfig { - int capacity; + int capacity; // logical cap in cells int n_layers; std::vector layers; - int kv_dtype; - int initial_capacity = 512; + int kv_dtype; // ET ScalarType the byte layer stores K/V in + int initial_capacity = 512; // starting pool size; grows lazily to capacity + // Max tokens per step; a ring layer sizes slots to window + max_write - 1. + // Unset = each ring layer uses its own window. std::optional max_write; }; -// Whether `cfg` satisfies the contract above. Callers must check this before -// constructing a cache: the `layers` broadcast rule is indexed directly, so a -// list that is neither size 1 nor n_layers reads past the end. Reported as a -// bool rather than thrown, so each backend picks its own failure mode. +// Whether `cfg` satisfies the contract above. inline bool valid(const CacheConfig& cfg) { - // initial_capacity may be 0 (allocate nothing up front) but not negative, and - // may exceed capacity -- the byte layer clamps it. + // initial_capacity may be 0 but not negative, and may exceed capacity -- the + // byte layer clamps it. return cfg.capacity > 0 && cfg.n_layers > 0 && cfg.initial_capacity >= 0 && (cfg.layers.size() == 1 || cfg.layers.size() == static_cast(cfg.n_layers)); diff --git a/extension/llm/cache/cache_registry.cpp b/extension/llm/cache/cache_registry.cpp index 225868e68cf..d54d05ddc73 100644 --- a/extension/llm/cache/cache_registry.cpp +++ b/extension/llm/cache/cache_registry.cpp @@ -69,6 +69,14 @@ Result> CacheBuilderRegistry::build( kind.c_str()); builder = it->second; } + // Checked here rather than in each cache: `layers` is indexed directly, so a + // list that is neither size 1 nor n_layers reads past the end. + ET_CHECK_OR_RETURN_ERROR( + valid(cfg), + InvalidArgument, + "cache: invalid CacheConfig for %s:%s", + backend_id.c_str(), + kind.c_str()); return builder(cfg); } diff --git a/extension/llm/cache/cache_registry.h b/extension/llm/cache/cache_registry.h index 1fbc438756d..ae730344c8a 100644 --- a/extension/llm/cache/cache_registry.h +++ b/extension/llm/cache/cache_registry.h @@ -12,9 +12,8 @@ // 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* 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. +// Caches are owned as CacheBase* and the faces are recovered through its as_* +// accessors (no RTTI), each null for a face the cache does not implement. #include #include diff --git a/extension/llm/cache/cell_cache.cpp b/extension/llm/cache/cell_cache.cpp new file mode 100644 index 00000000000..ceaa31a6686 --- /dev/null +++ b/extension/llm/cache/cell_cache.cpp @@ -0,0 +1,298 @@ +/* + * 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 { + +CellCache::CellCache(const CacheConfig& cfg) + : capacity_(cfg.capacity), + pos_(cfg.capacity, -1), + owners_(cfg.capacity, 0), + served_(cfg.n_layers, false) { + assert(valid(cfg)); + // One window per layer, from the same per-layer config the sequence cache + // reads. Layers agreeing on a window share a step. + windows_.reserve(cfg.n_layers); + for (int l = 0; l < cfg.n_layers; ++l) { + const LayerConfig& lc = + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; + windows_.push_back( + lc.policy.kind == LayerPolicy::Kind::Ring ? lc.policy.window : 0); + } +} + +// -- CacheControl ------------------------------------------------------------ + +bool CellCache::can_extend(int n) const { + return capacity_ - used_count_ >= n; +} + +int CellCache::capacity() const { + return capacity_; +} + +void CellCache::clear() { + std::fill(pos_.begin(), pos_.end(), -1); + std::fill(owners_.begin(), owners_.end(), 0); + info_.fill(SeqInfo{}); + used_end_ = 0; + used_count_ = 0; + reserved_ = 0; + declared_ = false; + step_seq_ids_.clear(); + step_pos_.clear(); + std::fill(served_.begin(), served_.end(), false); + invalidate_steps(); +} + +// -- BatchControl ------------------------------------------------------------ + +bool CellCache::declare_step(const std::vector& seq_ids) { + if (seq_ids.empty() || !can_extend(static_cast(seq_ids.size()))) { + return false; + } + for (int32_t seq_id : seq_ids) { + if (!valid_seq(seq_id) || !live(seq_id)) { + return false; + } + } + step_seq_ids_ = seq_ids; + declared_ = true; + invalidate_steps(); + std::fill(served_.begin(), served_.end(), false); + return true; +} + +bool CellCache::live(int32_t seq_id) const { + return (reserved_ & bit(seq_id)) != 0; +} + +std::optional CellCache::seq_new() { + for (int32_t seq_id = 0; seq_id < kMaxSeqs; ++seq_id) { + if (!live(seq_id)) { + reserved_ |= bit(seq_id); + return seq_id; + } + } + return std::nullopt; +} + +std::optional CellCache::seq_clone( + int32_t src, + std::optional upto) { + if (!valid_seq(src) || info_[src].count == 0) { + return std::nullopt; + } + const std::optional dst = seq_new(); + if (!dst) { + return std::nullopt; + } + const uint64_t src_bit = bit(src), dst_bit = bit(*dst); + for (int i = 0; i < used_end_; ++i) { + if ((owners_[i] & src_bit) && (!upto || pos_[i] < *upto)) { + owners_[i] |= dst_bit; + } + } + rescan(*dst); + invalidate_steps(); + return dst; +} + +bool CellCache::seq_rm(int32_t seq_id, int p0, std::optional p1) { + if (!valid_seq(seq_id)) { + return false; + } + const uint64_t b = bit(seq_id); + for (int i = 0; i < used_end_; ++i) { + if ((owners_[i] & b) && pos_[i] >= p0 && (!p1 || pos_[i] < *p1)) { + owners_[i] &= ~b; + if (owners_[i] == 0) { + pos_[i] = -1; + --used_count_; + } + } + } + while (used_end_ > 0 && pos_[used_end_ - 1] < 0) { + --used_end_; + } + rescan(seq_id); + if (info_[seq_id].count == 0) { + reserved_ &= ~bit(seq_id); // the last slot went, so the id is free again + } + invalidate_steps(); + return true; +} + +int CellCache::seq_len(int32_t seq_id) const { + return valid_seq(seq_id) ? info_[seq_id].count : 0; +} + +int CellCache::next_pos(int32_t seq_id) const { + return valid_seq(seq_id) ? info_[seq_id].max_pos + 1 : 0; +} + +int CellCache::free_cells() const { + return capacity_ - used_count_; +} + +int CellCache::used_end() const { + return used_end_; +} + +// -- CellStepper ------------------------------------------------------------- + +const CellStep* +CellCache::place_step(int layer, const int32_t* positions, int length) { + if (layer < 0 || layer >= static_cast(windows_.size()) || + served_[layer]) { + return nullptr; // out of range, or a forward that skipped declare_step + } + if (!placed_) { + if (!declared_ || length != static_cast(step_seq_ids_.size())) { + return nullptr; // no declaration, or a token count disagreeing with it + } + if (!extends(positions, length)) { + return nullptr; // nothing mutated yet, so the step can be re-placed + } + step_pos_.assign(positions, positions + length); + if (!place()) { + return nullptr; + } + declared_ = false; // one declaration, one placement + placed_ = true; + } + served_[layer] = true; + return &step_for(windows_[layer]); +} + +// -- internals --------------------------------------------------------------- + +uint64_t CellCache::bit(int32_t seq_id) { + return uint64_t{1} << seq_id; +} + +bool CellCache::valid_seq(int32_t seq_id) { + return seq_id >= 0 && seq_id < kMaxSeqs; +} + +bool CellCache::extends(const int32_t* positions, int length) const { + std::array newest{}; + for (int s = 0; s < kMaxSeqs; ++s) { + newest[s] = info_[s].max_pos; + } + for (int i = 0; i < length; ++i) { + const int32_t seq_id = step_seq_ids_[i]; + if (positions[i] <= newest[seq_id]) { + return false; + } + newest[seq_id] = positions[i]; + } + return true; +} + +int CellCache::lowest_free(int from) const { + for (int i = from; i < capacity_; ++i) { + if (pos_[i] < 0) { + return i; + } + } + return -1; +} + +void CellCache::invalidate_steps() { + placed_ = false; + steps_.clear(); +} + +void CellCache::rescan(int32_t seq_id) { + const uint64_t b = bit(seq_id); + SeqInfo info; + for (int i = 0; i < used_end_; ++i) { + if (owners_[i] & b) { + info.max_pos = std::max(info.max_pos, pos_[i]); + ++info.count; + } + } + info_[seq_id] = info; +} + +void CellCache::claim(int cell, int32_t pos, int32_t seq_id) { + pos_[cell] = pos; + owners_[cell] = bit(seq_id); + used_end_ = std::max(used_end_, cell + 1); + ++used_count_; + SeqInfo& info = info_[seq_id]; + info.max_pos = std::max(info.max_pos, pos); + ++info.count; +} + +bool CellCache::place() { + const int length = static_cast(step_pos_.size()); + cells_.resize(length); + // A free cell leaves every cell below it occupied, so the next scan resumes + // past it. + int from = 0; + for (int i = 0; i < length; ++i) { + const int cell = lowest_free(from); + if (cell < 0) { + return false; + } + cells_[i] = cell; + from = cell + 1; + } + for (int i = 0; i < length; ++i) { + claim(cells_[i], step_pos_[i], step_seq_ids_[i]); + } + return true; +} + +const CellStep& CellCache::step_for(int window) { + auto it = steps_.find(window); + if (it != steps_.end()) { + return it->second; + } + CellStep step; + step.length = static_cast(cells_.size()); + step.read_len = used_end_; + step.cells = cells_; + step.mask_bits = build_mask(window); + return steps_.emplace(window, std::move(step)).first->second; +} + +std::vector CellCache::build_mask(int window) const { + const int length = static_cast(cells_.size()); + std::vector bits( + static_cast(length) * static_cast(used_end_), 0); + for (int i = 0; i < length; ++i) { + const uint64_t tok_bit = bit(step_seq_ids_[i]); + const int32_t tok_pos = step_pos_[i]; + // A flat layer reaches back to the start; a windowed one to its window. + const int32_t oldest = window > 0 ? tok_pos - window + 1 : 0; + uint8_t* row = bits.data() + static_cast(i) * used_end_; + for (int j = 0; j < used_end_; ++j) { + row[j] = static_cast( + pos_[j] >= 0 && // occupied: a freed cell holds nothing + (owners_[j] & tok_bit) && // one of this query's sequences + pos_[j] <= tok_pos && // not the future; <= so a query sees itself + pos_[j] >= oldest); // within the window + } + } + return bits; +} + +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h new file mode 100644 index 00000000000..2c519d9d28c --- /dev/null +++ b/extension/llm/cache/cell_cache.h @@ -0,0 +1,164 @@ +/* + * 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 cell layout: many sequences over one pool of per-token cells. A cell +// holds one token's position and the sequences owning it, so a sequence needs +// no contiguous range and a fork sets a second owner bit. The step's tokens sit +// on one flat axis, sequence identity supplied out-of-band by declare_step. +// Tensor-free; the byte layer holds the pools and applies what it is given. + +#include +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace cache { + +// Integer-only handoff to the byte layer, covering the whole forward: a cell +// means the same token in every layer's pool. +struct CellStep { + int length; + int read_len; // the window is cells [0, read_len) + std::vector cells; // cell per query token + std::vector mask_bits; // [length, read_len], 1 = attend +}; + +// Backend face of the cell layout. The step's first layer places its tokens and +// every later layer reuses that placement. `layer` selects the window, which +// decides the kind and mask, so a step is per policy and memoized for the +// forward. The returned step is owned by the cache and valid until the next +// verb. nullptr = no declaration, a token count disagreeing with it, a position +// a sequence already holds, a layer out of range, or a layer served twice. +class CellStepper { + public: + virtual ~CellStepper() = default; + virtual const CellStep* + place_step(int layer, const int32_t* positions, int length) = 0; +}; + +class CellCache : public CacheBase, public BatchControl, public CellStepper { + public: + // One bit per sequence in the owner bitset. + static constexpr int kMaxSeqs = 64; + + // Precondition: valid(cfg). CacheBuilderRegistry::build enforces it for + // registry-created caches; direct construction must check first. + explicit CellCache(const CacheConfig& cfg); + + CacheBase* base() { + return this; + } + BatchControl* as_batch_control() override { + return this; + } + CellStepper* as_cell_stepper() override { + return this; + } + + // -- CacheControl ------------------------------------------------------ + + bool can_extend(int n = 1) const override; + int capacity() const override; + void clear() override; + + // -- BatchControl ------------------------------------------------------ + + bool declare_step(const std::vector& seq_ids) override; + std::optional seq_new() override; + std::optional seq_clone(int32_t src, std::optional upto) + override; + bool seq_rm(int32_t seq_id, int p0, std::optional p1) override; + int seq_len(int32_t seq_id) const override; + int next_pos(int32_t seq_id) const override; + + int free_cells() const; + int used_end() const; + + // -- CellStepper ------------------------------------------------------- + + const CellStep* place_step(int layer, const int32_t* positions, int length) + override; + + private: + struct SeqInfo { + int count = 0; + int max_pos = -1; + }; + + static uint64_t bit(int32_t seq_id); + // Live from the moment seq_new hands the id out until its last slot goes. + // The bit alone answers this because a sequence can only take slots under an + // id declare_step accepted, and only a live id is accepted. + bool live(int32_t seq_id) const; + static bool valid_seq(int32_t seq_id); + + // Every position must be newer than what that sequence already holds, or it + // would own two cells for one token. Two cells with the same pos and owner + // are indistinguishable, so a branch is its own sequence. + bool extends(const int32_t* positions, int length) const; + + // Placement policy: the lowest free cell, so freed cells refill before the + // extent grows. The choice moves only the read window's width and how often a + // step fuses; the mask keys off pos/owners, never the index. + int lowest_free(int from) const; + + // Drop the step's placement and the per-window steps built from it, after + // the table moves underneath them. + void invalidate_steps(); + + // Recompute a sequence's summary after the verbs move cells under it. + void rescan(int32_t seq_id); + + void claim(int cell, int32_t pos, int32_t seq_id); + + // Claim a cell per token, shared by every layer of the forward. False = the + // pool cannot supply them; no cell is claimed until every one is found, so a + // refusal leaves the table unchanged. + bool place(); + + // One window's step, built the first time a layer with this window asks and + // kept for the rest of the step. + const CellStep& step_for(int window); + + // Query i attends cell j iff j is occupied, shares a sequence with i, is no + // newer than i, and on a windowed layer no older than its window. The step's + // cells are already placed, so a query sees itself and its earlier tokens. + std::vector build_mask(int window) const; + + int capacity_; + std::vector pos_; // per cell; -1 = free + std::vector owners_; // per cell; owning-sequence bitset + int used_count_ = 0; // occupied cells, so admission stays O(1) + int used_end_ = 0; // every occupied cell is in [0, used_end) + std::array info_{}; + uint64_t reserved_ = 0; // ids handed out by seq_new + + std::vector step_seq_ids_; // set by declare_step + std::vector step_pos_; // set when the step is placed + std::vector cells_; // the step's placement, shared by every layer + std::vector served_; // layers this step has already answered + std::vector windows_; // per layer; 0 = keeps all history + // window -> step, memoized per forward. Node-based is required: a step + // handed to one layer must survive another layer's insert. + std::map steps_; + bool declared_ = false; + bool placed_ = false; +}; + +} // namespace cache +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h index 905fa0bfa8b..5af756e6e24 100644 --- a/extension/llm/cache/sequence_cache.h +++ b/extension/llm/cache/sequence_cache.h @@ -14,6 +14,7 @@ // flat/ring model (gemma4) stays coherent. Tensor-free / ET-independent. #include +#include #include #include #include @@ -97,6 +98,7 @@ class SequenceCache : public CacheBase, public: explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity), max_write_(cfg.max_write) { + assert(valid(cfg)); 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. diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index d8577a09bbe..e41e91a6089 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -9,19 +9,27 @@ #include #include #include +#include #include #include +#include +#include +#include #include #include #include +using executorch::extension::llm::cache::BatchControl; 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::CellCache; +using executorch::extension::llm::cache::CellStep; +using executorch::extension::llm::cache::CellStepper; using executorch::extension::llm::cache::LayerConfig; using executorch::extension::llm::cache::LayerPolicy; using executorch::extension::llm::cache::make_unique_key; @@ -204,6 +212,19 @@ TEST_F(CacheTest, BuilderBuildsRegisteredKindElseError) { EXPECT_EQ(cache.get()->as_control()->capacity(), 32); EXPECT_EQ(reg.build("TestBackend", "missing", cfg).error(), Error::NotFound); + + // A layers list that is neither size 1 nor n_layers would be indexed past + // the end, so build refuses it before the cache is constructed. + EXPECT_EQ( + reg.build("TestBackend", "seq", CacheConfig{32, 3, {}}).error(), + Error::InvalidArgument); + EXPECT_EQ( + reg.build( + "TestBackend", + "seq", + CacheConfig{32, 3, {flat_layer(), flat_layer()}}) + .error(), + Error::InvalidArgument); } TEST_F(CacheTest, SessionInstallsOnCtorErasesOnDtor) { @@ -233,3 +254,427 @@ TEST_F(CacheTest, EtAdapterMapsResultsAndCodes) { EXPECT_EQ(et::rewind(cache, 9), Error::InvalidArgument); // cannot grow EXPECT_EQ(et::rewind(cache, 1), Error::Ok); } + +// ---- Cell layout ----------------------------------------------------------- + +namespace { +// One sequence's tokens in a step: `length` of them from `start_pos` onward. +struct SeqTokens { + int32_t seq_id; + int32_t start_pos; + int length; +}; + +// The per-token arrays a step is made of. +struct StepArgs { + std::vector seq_ids; + std::vector positions; +}; + +// Lay a step's sequences on one token axis, building the two arrays together +// so they cannot fall out of alignment. +StepArgs flatten_step(std::initializer_list sequences) { + StepArgs args; + for (const SeqTokens& t : sequences) { + for (int i = 0; i < t.length; ++i) { + args.seq_ids.push_back(t.seq_id); + args.positions.push_back(t.start_pos + i); + } + } + return args; +} + +// A cell cache and the two faces a caller holds: the runner drives the verbs, +// the backend places each step. +struct Cells { + explicit Cells( + int capacity, + std::vector layers = {flat_layer(), flat_layer()}) + : cache(CacheConfig{capacity, static_cast(layers.size()), layers}), + ctl(cache.as_batch_control()), + stepper(cache.as_cell_stepper()) {} + + // Ids come from the cache, so a test names sequences by allocating them. + // Unlike the face's, this one unwraps and fails the test if none is free. + int32_t seq_new() { + const auto id = ctl->seq_new(); + EXPECT_TRUE(id) << "no free sequence id"; + return id ? *id : -1; + } + + // One layer of an already-declared step. + const CellStep* place(int layer, std::vector positions) { + return stepper->place_step( + layer, positions.data(), static_cast(positions.size())); + } + + // A whole single-layer step: declare it, then place it. A nullptr from here + // is the placement's refusal -- admission failing is a test failure. + const CellStep* step( + std::vector seq_ids, + std::vector positions) { + EXPECT_TRUE(ctl->declare_step(seq_ids)) << "the step was not admitted"; + return place(/*layer=*/0, positions); + } + + const CellStep* step(StepArgs args) { + return step(std::move(args.seq_ids), std::move(args.positions)); + } + + CellCache cache; + BatchControl* ctl; + CellStepper* stepper; +}; + +// The mask row for query `i`, as a string of '.' and '1'. +std::string row(const CellStep& step, int i) { + EXPECT_EQ( + step.mask_bits.size(), + static_cast(step.length) * static_cast(step.read_len)); + if (step.mask_bits.empty()) { + return {}; + } + std::string out(step.read_len, '.'); + for (int j = 0; j < step.read_len; ++j) { + out[j] = step.mask_bits[i * step.read_len + j] ? '1' : '.'; + } + return out; +} +} // namespace + +TEST_F(CacheTest, CellSingleSequenceAttendsItsOwnPrefix) { + Cells c(16); + const int32_t s0 = c.seq_new(); + // seq 0 places 4 tokens at 0..3 + const auto* prefill = c.step(flatten_step({{s0, 0, 4}})); + ASSERT_NE(prefill, nullptr); + EXPECT_EQ(prefill->read_len, 4); + EXPECT_EQ(prefill->cells, (std::vector{0, 1, 2, 3})); + EXPECT_EQ(row(*prefill, 0), "1..."); + EXPECT_EQ(row(*prefill, 3), "1111"); + + const auto* decode = c.step(flatten_step({{s0, 4, 1}})); // one more at 4 + ASSERT_NE(decode, nullptr); + EXPECT_EQ(decode->cells, (std::vector{4})); + EXPECT_EQ(row(*decode, 0), "11111"); // the whole history it owns +} + +TEST_F(CacheTest, CellSecondSequenceForcesAnExplicitMask) { + Cells c(16); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + c.step(flatten_step({{s0, 0, 4}})); // seq 0 places 4 tokens at 0..3 + + const auto* step = + // seq 1 places 2 at positions seq 0 also holds + c.step(flatten_step({{s1, 0, 2}})); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->read_len, 6); + EXPECT_EQ(step->cells, (std::vector{4, 5})); + // sequence 1 sees none of sequence 0's cells, though they share positions + EXPECT_EQ(row(*step, 0), "....1."); + EXPECT_EQ(row(*step, 1), "....11"); +} + +TEST_F(CacheTest, CellBatchedPrefillKeepsSequencesApart) { + Cells c(16); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + + // One step, two prefills of different lengths on a single token axis. + const auto* step = c.step(flatten_step({{s0, 0, 3}, {s1, 0, 2}})); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->cells, (std::vector{0, 1, 2, 3, 4})); + EXPECT_EQ(step->read_len, 5); + + EXPECT_EQ(row(*step, 0), "1...."); // seq 0 causally over its own three + EXPECT_EQ(row(*step, 1), "11..."); + EXPECT_EQ(row(*step, 2), "111.."); + EXPECT_EQ(row(*step, 3), "...1."); // seq 1 sees none of seq 0's, though + EXPECT_EQ(row(*step, 4), "...11"); // they hold the same positions +} + +TEST_F(CacheTest, CellDecodeAndPrefillShareOneStep) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 2}})); // an existing conversation at 0..1 + + // Continuous batching: s0 decodes while s1 arrives and prefills. + const int32_t s1 = c.seq_new(); + const auto* step = c.step(flatten_step({{s0, 2, 1}, {s1, 0, 2}})); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->cells, (std::vector{2, 3, 4})); + + EXPECT_EQ(row(*step, 0), "111.."); // s0's decode sees its own history + EXPECT_EQ(row(*step, 1), "...1."); // s1 starts from nothing + EXPECT_EQ(row(*step, 2), "...11"); + EXPECT_EQ(c.ctl->next_pos(s0), 3); + EXPECT_EQ(c.ctl->next_pos(s1), 2); +} + +TEST_F(CacheTest, CellExtendsIsCheckedPerSequence) { + Cells c(16); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + c.step(flatten_step({{s0, 0, 1}, {s1, 0, 1}})); // both at position 0 + + // One sequence advancing does not license another to repeat: the check is + // against what each sequence itself already holds. + EXPECT_EQ(c.step({s0, s1}, {1, 0}), nullptr); + EXPECT_NE(c.step(flatten_step({{s0, 1, 1}, {s1, 1, 1}})), nullptr); +} + +TEST_F(CacheTest, CellPlacementIsSharedByEveryLayerOfTheStep) { + Cells c(16); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + const auto args = flatten_step({{s0, 0, 1}, {s1, 0, 1}}); + ASSERT_TRUE(c.ctl->declare_step(args.seq_ids)); + + const auto* first = c.place(0, args.positions); // layer 0 places the cells + ASSERT_NE(first, nullptr); + EXPECT_EQ(c.place(1, args.positions), first); // later layers reuse them + EXPECT_EQ(c.place(0, args.positions), nullptr); // asking twice is a new step + EXPECT_EQ(c.cache.free_cells(), 14); // placed once, not once per layer +} + +TEST_F(CacheTest, CellForkSharesCellsAndEvictionRefcounts) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 4}})); // seq 0 places 4 tokens at 0..3 + + ASSERT_EQ(c.cache.free_cells(), 12); // the four it placed, out of sixteen + + const auto s1 = c.ctl->seq_clone(s0, std::nullopt); + ASSERT_TRUE(s1); + EXPECT_EQ(c.cache.free_cells(), 12); // no cell, no byte copied + EXPECT_EQ(c.ctl->seq_len(*s1), 4); + EXPECT_EQ(c.ctl->next_pos(*s1), 4); + + c.ctl->seq_rm(s0, 0, std::nullopt); + EXPECT_EQ(c.ctl->seq_len(s0), 0); + EXPECT_EQ(c.ctl->seq_len(*s1), 4); // the fork still owns them + EXPECT_EQ(c.cache.free_cells(), 12); // so nothing is reclaimed yet + + c.ctl->seq_rm(*s1, 0, std::nullopt); + EXPECT_EQ(c.cache.free_cells(), 16); + EXPECT_EQ(c.cache.used_end(), 0); // the extent comes back too +} + +TEST_F(CacheTest, CellSeqNewHandsOutIdsUntilTheyAreReleased) { + Cells c(16); + const auto a = c.ctl->seq_new(); + const auto b = c.ctl->seq_new(); + ASSERT_TRUE(a && b); + EXPECT_NE(*a, *b); + + // Reserved before it holds anything, so the next call cannot hand it out. + EXPECT_EQ(c.ctl->seq_len(*a), 0); + EXPECT_NE(*c.ctl->seq_new(), *a); + + // Removing everything a sequence holds returns its id, whether or not it + // ever held a slot. + ASSERT_TRUE(c.ctl->seq_rm(*a, 0, std::nullopt)); + EXPECT_EQ(*c.ctl->seq_new(), *a); + + // An id nobody was handed does not start a sequence. + EXPECT_FALSE(c.ctl->declare_step({40})); + + while (c.ctl->seq_new()) { + } + EXPECT_FALSE(c.ctl->seq_new()); // every id is now in use + + c.ctl->clear(); + EXPECT_EQ(*c.ctl->seq_new(), 0); +} + +TEST_F(CacheTest, CellForkCanShareAPrefixOnly) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 4}})); // positions 0..3 + + const auto s1 = c.ctl->seq_clone(s0, /*upto=*/2); // positions 0..1 only + ASSERT_TRUE(s1); + EXPECT_EQ(c.ctl->seq_len(*s1), 2); + EXPECT_EQ(c.ctl->next_pos(*s1), 2); // the fork resumes where the prefix ends + EXPECT_EQ(c.ctl->seq_len(s0), 4); // the source keeps all of its own + EXPECT_EQ(c.cache.free_cells(), 12); // still no cell copied + + // Removing the source's shared range frees nothing: the fork still owns it. + ASSERT_TRUE(c.ctl->seq_rm(s0, 0, 2)); + EXPECT_EQ(c.cache.free_cells(), 12); + EXPECT_EQ(c.ctl->seq_len(*s1), 2); +} + +TEST_F(CacheTest, CellWindowAndSequenceBothNarrowTheMask) { + // The two mask bounds together: a query sees only its own sequence, and only + // inside its window. + Cells c(16, {ring_layer(2), ring_layer(2)}); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + c.step(flatten_step({{s0, 0, 3}})); // seq 0 at 0..2 + const auto* step = c.step(flatten_step({{s1, 0, 1}, {s0, 3, 1}})); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->read_len, 5); + + // Seq 1 holds only its own new cell, and its window reaches no further. + EXPECT_EQ(row(*step, 0), "...1."); + // Seq 0 at position 3 sees positions 2 and 3 -- cells 2 and 4 -- but not its + // own cells 0 and 1, which the window excludes. + EXPECT_EQ(row(*step, 1), "..1.1"); +} + +TEST_F(CacheTest, CellVerbBetweenLayersInvalidatesTheStep) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 2}})); + + ASSERT_TRUE(c.ctl->declare_step({s0})); + ASSERT_NE(c.place(0, {2}), nullptr); + // A verb rebuilds the table under the step: the placement no longer stands + // and the remaining layers are refused. + ASSERT_TRUE(c.ctl->seq_rm(s0, 0, 1)); + EXPECT_EQ(c.place(1, {2}), nullptr); +} + +TEST_F(CacheTest, CellRangedRemovalFreesOnlyThatWindow) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 5}})); // seq 0 places 5 tokens at 0..4 + + c.ctl->seq_rm(s0, 0, 2); // sliding window: drop the oldest two + EXPECT_EQ(c.ctl->seq_len(s0), 3); + EXPECT_EQ(c.cache.free_cells(), 13); + EXPECT_EQ(c.ctl->next_pos(s0), 5); // a count is not a position + + c.ctl->seq_rm(s0, 4, std::nullopt); // backtrack: drop position 4 onwards + EXPECT_EQ(c.ctl->seq_len(s0), 2); + EXPECT_EQ(c.ctl->next_pos(s0), 4); +} + +TEST_F(CacheTest, CellRefillingHolesKeepsTheMaskOnPositions) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 4}})); // seq 0 places 4 tokens at 0..3 + c.ctl->seq_rm(s0, 0, 2); // free the oldest cells, leaving holes at 0 and 1 + + // The new tokens take those holes, so the sequence's cells no longer ascend + // with its positions -- the mask keys off pos/owners, never the index. + // seq 0 places two more at 4..5 + const auto* step = c.step(flatten_step({{s0, 4, 2}})); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->cells, (std::vector{0, 1})); + // cells 0,1 hold positions 4,5; cells 2,3 hold 2,3 + EXPECT_EQ(row(*step, 0), "1.11"); // query at 4 sees 2, 3 and itself + EXPECT_EQ(row(*step, 1), "1111"); // query at 5 sees all of them +} + +TEST_F(CacheTest, CellRejectsAPositionASequenceStillHolds) { + Cells c(16); + const int32_t s0 = c.seq_new(); + c.step(flatten_step({{s0, 0, 4}})); // seq 0 places 4 tokens at 0..3 + c.ctl->seq_rm(s0, 3, std::nullopt); // free the newest cell + + // A step only extends its sequences. Writing 0 again would leave seq 0 with + // two cells for one token and a window that is no longer a causal prefix. + EXPECT_EQ(c.step({s0}, {0}), nullptr); + // nor may its own tokens descend + EXPECT_EQ(c.step({s0, s0}, {5, 4}), nullptr); + EXPECT_NE(c.step({s0}, {3}), nullptr); // the position it just freed is fine + + // A refusal places nothing, so the declaration stands and the same step can + // be placed again with corrected positions. + ASSERT_TRUE(c.ctl->declare_step({s0})); + EXPECT_EQ(c.place(0, {0}), nullptr); + EXPECT_NE(c.place(0, {4}), nullptr); +} + +TEST_F(CacheTest, CellWindowBoundsEachQueryByPosition) { + { + Cells c(16, {ring_layer(4)}); + // span 3 < window 4: the last query still sees position 0. + EXPECT_EQ(row(*c.step(flatten_step({{c.seq_new(), 0, 4}})), 3), "1111"); + } + { + Cells c(16, {ring_layer(3)}); + // span 3 == window 3: position 0 is one too old for the query at 3. + EXPECT_EQ(row(*c.step(flatten_step({{c.seq_new(), 0, 4}})), 3), ".111"); + } + { + Cells c(16, {ring_layer(2)}); + // Two cells, five positions apart: the span is what counts. + const int32_t s = c.seq_new(); + c.step({s}, {0}); + EXPECT_EQ(row(*c.step({s}, {5}), 0), ".1"); + } + { + Cells c(16, {ring_layer(8)}); + // Sparse but inside the window, so both cells stay visible. + const int32_t s = c.seq_new(); + c.step({s}, {0}); + EXPECT_EQ(row(*c.step({s}, {5}), 0), "11"); + } +} + +TEST_F(CacheTest, CellLayersSharingAWindowShareAStep) { + // gemma-style: one flat layer beside two windowed ones. The placement is + // shared, the step is per policy -- so the flat layer keeps the whole prefix + // while the windowed pair gets one banded step between them. + Cells c(16, {flat_layer(), ring_layer(2), ring_layer(2)}); + const int32_t s0 = c.seq_new(); + const auto args = flatten_step({{s0, 0, 4}}); + ASSERT_TRUE(c.ctl->declare_step(args.seq_ids)); + + const auto* flat = c.place(0, args.positions); + const auto* windowed = c.place(1, args.positions); + ASSERT_NE(flat, nullptr); + ASSERT_NE(windowed, nullptr); + EXPECT_EQ(row(*flat, 2), "111."); // the flat layer keeps position 0 + EXPECT_EQ(row(*windowed, 2), ".11."); // the windowed one drops it + + EXPECT_EQ(c.place(2, args.positions), windowed); // same window, same step + EXPECT_EQ(c.cache.free_cells(), 12); // placed once for all three layers +} + +TEST_F(CacheTest, CellStepProtocolIsEnforced) { + Cells c(4); + const int32_t s0 = c.seq_new(); + const std::vector seqs{s0, s0}; + EXPECT_FALSE(c.ctl->declare_step({})); // no tokens + EXPECT_FALSE(c.ctl->declare_step({s0, s0, s0, s0, s0})); // 5 tokens, 4 cells + EXPECT_FALSE( + c.ctl->declare_step({CellCache::kMaxSeqs})); // seq id past the last bit + // The verbs report a bad sequence rather than doing nothing quietly. + EXPECT_FALSE(c.ctl->seq_clone(CellCache::kMaxSeqs, std::nullopt)); + EXPECT_FALSE(c.ctl->seq_clone(s0 + 30, std::nullopt)); // src holds nothing + EXPECT_FALSE(c.ctl->seq_rm(-1, 0, std::nullopt)); + + EXPECT_EQ(c.place(0, {0, 1}), nullptr); // no declaration + ASSERT_TRUE(c.ctl->declare_step(seqs)); + EXPECT_EQ(c.place(0, {0}), nullptr); // token count disagrees + ASSERT_NE(c.place(0, {0, 1}), nullptr); // 2 tokens, as declared + + EXPECT_EQ(c.place(0, {2, 3}), nullptr); // a second step, no declare_step + + EXPECT_EQ(c.place(-1, {0, 1}), nullptr); // layers outside the model + EXPECT_EQ(c.place(2, {0, 1}), nullptr); +} + +TEST_F(CacheTest, CellClearReturnsEveryCell) { + Cells c(4); + const int32_t s0 = c.seq_new(); + EXPECT_EQ(c.cache.capacity(), 4); + EXPECT_TRUE(c.ctl->can_extend(4)); + + c.step(flatten_step({{s0, 0, 3}})); + EXPECT_FALSE(c.ctl->can_extend(2)); // one cell left + EXPECT_EQ(c.ctl->seq_len(s0), 3); + + c.ctl->clear(); + EXPECT_TRUE(c.ctl->can_extend(4)); + EXPECT_EQ(c.cache.free_cells(), 4); + EXPECT_EQ(c.cache.used_end(), 0); + EXPECT_EQ(c.ctl->seq_len(s0), 0); + EXPECT_EQ(c.ctl->next_pos(s0), 0); // the sequence is gone + EXPECT_EQ(c.place(0, {0}), nullptr); // and the step went with it +}