diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index a7bcfd66b3b..3ab424830e0 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -13,9 +13,10 @@ // 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 +// 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. #include @@ -58,10 +59,10 @@ struct Run { }; // 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. +// 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. struct SeqStepPlan { Run write[2]; int n_write; @@ -84,26 +85,27 @@ class SequencePlanner { virtual void commit(const SeqStepPlan& plan) = 0; }; -// Per-layer layout behavior (e.g. full-history flat). Pure: plan() has no side -// effects, so the controller (SequenceCache) owns length. +// Per-layer layout behavior (flat = full history; ring = sliding window). 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. + // Oldest logical position still retained given the current length (0 for + // flat; length - window for ring). 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 + enum class Kind : int { + Flat = 0, + Ring = 1 + }; // serialized values: append-only Kind kind = Kind::Flat; - int window = - 0; // sliding-window size for windowing policies; 0 = full history + int window = 0; // Ring: sliding-window size; must be 0 when Flat }; // Per-layer architecture facts + cache policy. @@ -114,13 +116,17 @@ struct LayerConfig { }; // 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. +// the logical cap; 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. struct CacheConfig { int capacity; int n_layers; std::vector layers; int initial_capacity = 512; + std::optional max_write; }; } // namespace cache diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h index c9886290073..905fa0bfa8b 100644 --- a/extension/llm/cache/sequence_cache.h +++ b/extension/llm/cache/sequence_cache.h @@ -8,10 +8,10 @@ #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. +// The neutral single-sequence controller (SequenceCache) and its layout +// policies (FlatPolicy / RingPolicy). SequenceCache owns the one logical length +// for the whole model and dispatches per-layer layout to a policy, so a mixed +// flat/ring model (gemma4) stays coherent. Tensor-free / ET-independent. #include #include @@ -43,14 +43,60 @@ class FlatPolicy final : public LayoutPolicy { } }; +// Sliding window of `window` cells over a ring of `window + max_write - 1` +// slots. The ring is oversized so a step of up to max_write tokens fits without +// overwriting cells earlier queries in the same step still attend to; the +// backend masks each query to its own window within the read span. +class RingPolicy final : public LayoutPolicy { + public: + RingPolicy(int window, int max_write) + : window_(window), ring_size_(window + max_write - 1) {} + + int retained_from(int length) const override { + return length > window_ ? length - window_ : 0; + } + SeqStepPlan plan(int position, int T) const override { + const int end = position + T; + SeqStepPlan p{}; + // Write this step's T cells (T <= max_write), wrapping the ring if needed. + p.n_write = split_runs(position, T, p.write); + // Read the union of the step's per-query windows, [position - window + 1, + // end), in logical (oldest -> newest) order. + const int rstart = position - window_ + 1 > 0 ? position - window_ + 1 : 0; + const int rlen = end - rstart; + p.n_read = split_runs(rstart, rlen, p.read); + p.read_base_pos = rstart; + return p; + } + + private: + // Split a logical run [start, start+len) into up to two physical runs in the + // ring (wrapping at ring_size_). Precondition: len <= ring_size_. + int split_runs(int start, int len, Run out[2]) const { + const int phys_start = start % ring_size_; + const int first_len = std::min(len, ring_size_ - phys_start); + out[0] = Run{phys_start, first_len}; + if (first_len < len) { + out[1] = Run{0, len - first_len}; + return 2; + } + return 1; + } + + int window_; + int ring_size_; +}; + // 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. +// deduped by (kind, window), so a uniform or two-kind (gemma4) model holds one +// or two policy objects. class SequenceCache : public CacheBase, public SequenceControl, public SequencePlanner { public: - explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity) { + explicit SequenceCache(const CacheConfig& cfg) + : capacity_(cfg.capacity), max_write_(cfg.max_write) { 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. @@ -105,6 +151,13 @@ class SequenceCache : public CacheBase, if (position + T > capacity_) { return std::nullopt; } + // A ring layer's slots are sized window + max_write - 1 (max_write defaults + // to the window), so a step larger than that would overrun it. + const LayerPolicy& spec = specs_[layer_to_policy_[layer]]; + if (spec.kind == LayerPolicy::Kind::Ring && + T > (max_write_ ? *max_write_ : spec.window)) { + return std::nullopt; + } return policies_[layer_to_policy_[layer]]->plan(position, T); } void commit(const SeqStepPlan& plan) override { @@ -129,11 +182,17 @@ class SequenceCache : public CacheBase, policies_.push_back(make_policy(lp)); return static_cast(policies_.size() - 1); } - std::unique_ptr make_policy(const LayerPolicy&) const { + std::unique_ptr make_policy(const LayerPolicy& lp) const { + if (lp.kind == LayerPolicy::Kind::Ring) { + // Unset max_write -> default a ring layer to its own window (ring 2w-1). + const int mw = max_write_ ? *max_write_ : lp.window; + return std::make_unique(lp.window, mw); + } return std::make_unique(); } int capacity_; + std::optional max_write_; int length_ = 0; std::vector specs_; // parallel to policies_, for dedup std::vector> policies_; diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index 9d69c1ea739..d8577a09bbe 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -33,6 +33,9 @@ namespace { LayerConfig flat_layer() { return LayerConfig{LayerPolicy{LayerPolicy::Kind::Flat, 0}, 2, 8}; } +LayerConfig ring_layer(int window) { + return LayerConfig{LayerPolicy{LayerPolicy::Kind::Ring, window}, 2, 8}; +} } // namespace // Initializes the ExecuTorch PAL so the ET adapter's error paths (which ET_LOG) @@ -66,6 +69,58 @@ TEST_F(CacheTest, FlatPlanAppendsAndReadsAllHistory) { EXPECT_EQ(p1->read[0].len, 5); } +// ---- Ring policy: wrap, eviction, read base position ----------------------- + +TEST_F(CacheTest, RingPlanWrapsAndEvicts) { + // window 4, max_write unset -> defaults to window -> ring of 2*4 - 1 = 7. + SequenceCache cache(CacheConfig{100, 1, {ring_layer(4)}}); + + // Prefill chunk of 4 at position 0: one write run, reads [0,4), no wrap. + auto p0 = cache.plan(0, 0, 4); + 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); + cache.commit(*p0); + + // Decode at position 7: write wraps to slot 0 (7 % 7); the read window + // [4,8) wraps the ring -> phys [4,6] then [0,0]. Positions 0-3 evicted. + auto p1 = cache.plan(0, 7, 1); + ASSERT_TRUE(p1.has_value()); + EXPECT_EQ(p1->n_write, 1); + EXPECT_EQ(p1->write[0].start, 0); + EXPECT_EQ(p1->write[0].len, 1); + EXPECT_EQ(p1->n_read, 2); + EXPECT_EQ(p1->read[0].start, 4); + EXPECT_EQ(p1->read[0].len, 3); + EXPECT_EQ(p1->read[1].start, 0); + EXPECT_EQ(p1->read[1].len, 1); + EXPECT_EQ(p1->read_base_pos, 4); // oldest retained is logical position 4 + + // A step larger than max_write would overrun the ring -> rejected. + EXPECT_FALSE(cache.plan(0, 0, 5).has_value()); +} + +// ---- Mixed flat/ring: one shared length across layers ---------------------- + +TEST_F(CacheTest, MixedFlatRingShareOneLength) { + SequenceCache cache(CacheConfig{100, 2, {flat_layer(), ring_layer(4)}}); + + // Same step drives both layers (T <= window); commit once. + auto p = cache.plan(0, 0, 3); + ASSERT_TRUE(p.has_value()); + ASSERT_TRUE(cache.plan(1, 0, 3).has_value()); + cache.commit(*p); + + // Length advanced once (not per layer): length == 3. + EXPECT_TRUE(cache.can_extend(97)); + EXPECT_FALSE(cache.can_extend(98)); +} + // ---- Admission / rewind ---------------------------------------------------- TEST_F(CacheTest, CanExtendBoundedByCapacity) { @@ -89,6 +144,21 @@ TEST_F(CacheTest, FlatRewindsFreelyToZero) { EXPECT_TRUE(cache.can_extend(8)); } +TEST_F(CacheTest, RewindBoundedByRingWindow) { + SequenceCache cache(CacheConfig{100, 2, {flat_layer(), ring_layer(4)}}); + // Advance length to 10 in chunks of 2 (<= window). + for (int pos = 0; pos < 10; pos += 2) { + auto p = cache.plan(0, pos, 2); + ASSERT_TRUE(p.has_value()); + ASSERT_TRUE(cache.plan(1, pos, 2).has_value()); + cache.commit(*p); + } + // ring(4) has evicted everything older than 10 - 4 = 6. + EXPECT_FALSE(cache.rewind(5)); // older than the ring retains + EXPECT_TRUE(cache.rewind(6)); // exactly the floor + EXPECT_FALSE(cache.rewind(11)); // cannot grow +} + // ---- Faces / registry / session -------------------------------------------- TEST_F(CacheTest, FaceRecoveryReturnsSameObject) {