diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index ea88e5ed686..e1dc76af20e 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -875,7 +875,7 @@ class MLXSequenceCache final : public cache::Cache, int capacity() const override { return capacity_; } - int length() const { + int length() const override { return length_; } // Includes persistent growth in a partially constructed forward. diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 09fab2fc0c6..60dbefe0933 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -90,6 +90,9 @@ class ET_EXPERIMENTAL SequenceControl : public CacheControl { // Room for n more tokens within capacity. virtual bool can_extend(int n = 1) const = 0; + // Where the sequence stands: one past its newest position, and so where it + // continues. Positions are dense from 0, so this is also how many it holds. + virtual int length() const = 0; // Keep [0, position), dropping the rest. False = a target it has not // reached, or one older than an evicting layer still retains. virtual bool rewind(int position) = 0; diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h index eeea8e0c065..6dc12b1d8d6 100644 --- a/extension/llm/cache/sequence_cache.h +++ b/extension/llm/cache/sequence_cache.h @@ -189,9 +189,7 @@ class ET_EXPERIMENTAL SequenceCache : public Cache, length_ = 0; written_ = 0; } - // Positions are dense from 0, so this is both what the sequence holds and - // where its next token goes. - int length() const { + int length() const override { return length_; } // Whether rewind(position) would be accepted, without moving the length. diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index c073428d1ed..e1df5fc1221 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -244,6 +244,21 @@ TEST_F(CacheTest, FaceRecoveryReturnsSameObject) { EXPECT_EQ(plan->read[0].len, 1); } +TEST_F(CacheTest, ControlFaceReportsTheLength) { + SequenceCache cache(CacheGeometry{{flat_layer()}}, config(8)); + SequenceControl* control = static_cast(cache).as(); + ASSERT_NE(control, nullptr); + EXPECT_EQ(control->length(), 0); + auto plan = cache.plan(0, 0, 3); + ASSERT_TRUE(plan.has_value()); + cache.commit(*plan); + EXPECT_EQ(control->length(), 3); + ASSERT_TRUE(control->rewind(1)); + EXPECT_EQ(control->length(), 1); + control->clear(); + EXPECT_EQ(control->length(), 0); +} + TEST_F(CacheTest, NullFaceIdsNeverMatch) { using executorch::extension::llm::cache::same_face; EXPECT_FALSE(same_face(nullptr, nullptr));