From 38e6b21b0d5e858031c2c1209015fabea7919945 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Mon, 17 Aug 2026 15:18:36 -0700 Subject: [PATCH] Widen the MLX cache face to per-token positions and lift out the pool --- backends/mlx/runtime/MLXCache.h | 10 +- backends/mlx/runtime/MLXInterpreter.h | 29 +++- backends/mlx/runtime/MLXPool.h | 113 ++++++++++++++++ backends/mlx/runtime/MLXSequenceCache.h | 126 ++++-------------- backends/mlx/test/mlx_sequence_cache_test.cpp | 79 ++++++++--- 5 files changed, 228 insertions(+), 129 deletions(-) create mode 100644 backends/mlx/runtime/MLXPool.h diff --git a/backends/mlx/runtime/MLXCache.h b/backends/mlx/runtime/MLXCache.h index 98f144cfdf9..2fb1975bfa7 100644 --- a/backends/mlx/runtime/MLXCache.h +++ b/backends/mlx/runtime/MLXCache.h @@ -36,14 +36,12 @@ class MLXCache { public: virtual ~MLXCache() = default; - // Write this step's K/V for `layer` at `position` (the run's logical start); - // return the window + mask kind. k/v are BHSD. `position` is a host int -- - // the caller reads it off the graph so the cache stays pure graph + integer - // bookkeeping. The cache owns the mask: a multi-token chain is Causal, a - // single decode token is None. + // Write this step's K/V for `layer` at `positions`, one host int per query + // token, and return the window plus the mask kind. k/v are BHSD. virtual AttendSpec update_and_fetch( int layer, - int position, + const int32_t* positions, + int length, const Tensor& k, const Tensor& v, StreamOrDevice s) = 0; diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index b14fdf6903c..5f8cfc60519 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -11,6 +11,9 @@ #include "MLXCache.h" #include "MLXExecutor.h" +#include +#include + #include #include #include @@ -310,20 +313,33 @@ inline void exec_update_and_attend( // The cache does the KV write + read and declares the mask; the handler owns // the query side (q, scale) and calls SDPA. const array& q = st.const_tensor_ref(n.q); - // The run's start is position[0], read host-side so the cache stays pure - // graph + integer bookkeeping. Every layer of a step reads the same position + // One position per query token, read host-side so the cache stays pure graph + // + integer bookkeeping. Every layer of a step reads the same position // tensor, so evaluating it in place costs one sync for the first layer and // nothing for the rest -- casting first would instead build a fresh array per // layer and sync on each one. auto pos = st.const_tensor_ref(n.position); eval(pos); - int position; + const int length = static_cast(pos.size()); + if (length != static_cast(st.const_tensor_ref(n.k).shape(2))) { + throw std::runtime_error( + "update_and_attend: position must hold one entry per query token"); + } + // int32 is passed straight through; only an int64 input needs narrowing. + std::vector int32_positions; + const int32_t* positions = nullptr; switch (pos.dtype()) { case ::mlx::core::int32: - position = pos.data()[0]; + positions = pos.data(); break; case ::mlx::core::int64: - position = static_cast(pos.data()[0]); + int32_positions.resize(static_cast(length)); + std::transform( + pos.data(), + pos.data() + length, + int32_positions.begin(), + [](int64_t p) { return static_cast(p); }); + positions = int32_positions.data(); break; default: throw std::runtime_error( @@ -332,7 +348,8 @@ inline void exec_update_and_attend( } AttendSpec spec = st.cache->update_and_fetch( *n.layer_id, - position, + positions, + length, st.const_tensor_ref(n.k), st.const_tensor_ref(n.v), s); diff --git a/backends/mlx/runtime/MLXPool.h b/backends/mlx/runtime/MLXPool.h new file mode 100644 index 00000000000..ba3622ce5e5 --- /dev/null +++ b/backends/mlx/runtime/MLXPool.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 + +#include +#include +#include + +#include "MLXExecutor.h" // Tensor, StreamOrDevice + +namespace executorch { +namespace backends { +namespace mlx { + +// Per-layer K or V store, SDPA-major [1, H, slots, D] (cells on axis 2). The +// caller hands down physical slot ranges (it has already applied any ring +// modulo), so the pool is layout-agnostic: policies differ only in how many +// slots the layer asks for and how many ranges a step produces. +class Pool { + public: + // initial_slots above max_slots is clamped, not rejected: the config default + // exceeds the cap of any smaller cache, so this is the normal path. + Pool(int initial_slots, int max_slots, int H, int D, ::mlx::core::Dtype dtype) + : dtype_(dtype), + max_slots_(max_slots), + buf_(::mlx::core::zeros( + ::mlx::core::Shape{1, H, std::min(initial_slots, max_slots), D}, + dtype)) {} + + // Place `update` at slot `start`, casting to the storage dtype if it differs. + void write(int start, int len, const Tensor& update, StreamOrDevice s) { + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + if (start < 0 || start + len > max_slots_) { + throw std::runtime_error("Pool::write: run out of bounds"); + } + if (static_cast(update.shape(2)) != len) { + throw std::runtime_error("Pool::write: update length != run length"); + } + if (static_cast(update.shape(1)) != H || + static_cast(update.shape(3)) != D) { + throw std::runtime_error("Pool::write: K/V heads/dim mismatch"); + } + maybe_grow(start + len, s); + const Tensor u = update.dtype() == dtype_ + ? update + : ::mlx::core::astype(update, dtype_, s); + buf_ = ::mlx::core::slice_update( + buf_, + u, + ::mlx::core::Shape{0, 0, start, 0}, + ::mlx::core::Shape{1, H, start + len, D}, + s); + } + + // Slots [start, start+len). A ring read starts mid-pool, so the start matters + // here as much as it does for a write. + Tensor read(int start, int len, StreamOrDevice s) const { + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + if (start < 0 || start + len > slots()) { + throw std::runtime_error("Pool::read: run out of bounds"); + } + return ::mlx::core::slice( + buf_, + ::mlx::core::Shape{0, 0, start, 0}, + ::mlx::core::Shape{1, H, start + len, D}, + ::mlx::core::Shape{1, 1, 1, 1}, + s); + } + + // Slots currently allocated; grows toward max_slots on demand. + int slots() const { + return static_cast(buf_.shape(2)); + } + + private: + // Make room for `needed` slots, growing only if the pool is short: double + // until it fits, never past max_slots_. Cells keep their index, so growth is + // a zero-pad on the cell axis. + void maybe_grow(int needed, StreamOrDevice s) { + const int cur = slots(); + if (needed <= cur) { + return; + } + int next = std::max(cur, 1); // an empty pool has nothing to double + while (next < needed) { + next *= 2; + } + // The last doubling can overshoot; write() already bounds `needed` by + // max_slots_, so clamping here cannot undershoot it. + next = std::min(next, max_slots_); + const int H = static_cast(buf_.shape(1)); + const int D = static_cast(buf_.shape(3)); + Tensor pad = + ::mlx::core::zeros(::mlx::core::Shape{1, H, next - cur, D}, dtype_); + buf_ = ::mlx::core::concatenate(std::vector{buf_, pad}, 2, s); + } + + ::mlx::core::Dtype dtype_; + int max_slots_; + Tensor buf_; +}; + +} // namespace mlx +} // namespace backends +} // namespace executorch diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 9e644c3d977..2837a58f813 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -8,7 +8,6 @@ #pragma once -#include #include #include #include @@ -16,6 +15,7 @@ #include "MLXCache.h" // AttendSpec, MLXCache #include "MLXExecutor.h" // resolve_dtype +#include "MLXPool.h" // Pool #include #include @@ -26,97 +26,6 @@ namespace mlx { namespace cache = ::executorch::extension::llm::cache; -// Per-layer K or V store, SDPA-major [1, H, slots, D] (cells on axis 2). The -// planner hands down physical runs (it has already applied any ring modulo), so -// the pool is layout-agnostic: flat and ring differ only in how many slots the -// layer asks for and how many runs a step produces. -class Pool { - public: - // initial_slots above max_slots is clamped, not rejected: the config default - // exceeds the cap of any smaller cache, so this is the normal path. - Pool(int initial_slots, int max_slots, int H, int D, ::mlx::core::Dtype dtype) - : dtype_(dtype), - max_slots_(max_slots), - buf_(::mlx::core::zeros( - ::mlx::core::Shape{1, H, std::min(initial_slots, max_slots), D}, - dtype)) {} - - // Place `update` at the run's physical start, casting to the storage dtype if - // it differs. - void write(const cache::Run& run, const Tensor& update, StreamOrDevice s) { - const int H = static_cast(buf_.shape(1)); - const int D = static_cast(buf_.shape(3)); - if (run.start < 0 || run.start + run.len > max_slots_) { - throw std::runtime_error("Pool::write: run out of bounds"); - } - if (static_cast(update.shape(2)) != run.len) { - throw std::runtime_error("Pool::write: update length != run length"); - } - if (static_cast(update.shape(1)) != H || - static_cast(update.shape(3)) != D) { - throw std::runtime_error("Pool::write: K/V heads/dim mismatch"); - } - maybe_grow(run.start + run.len, s); - const Tensor u = update.dtype() == dtype_ - ? update - : ::mlx::core::astype(update, dtype_, s); - buf_ = ::mlx::core::slice_update( - buf_, - u, - ::mlx::core::Shape{0, 0, run.start, 0}, - ::mlx::core::Shape{1, H, run.start + run.len, D}, - s); - } - - // The run's cells, [start, start+len). Ring reads start mid-pool, so the run - // start matters here as much as it does for a write. - Tensor read(const cache::Run& run, StreamOrDevice s) const { - const int H = static_cast(buf_.shape(1)); - const int D = static_cast(buf_.shape(3)); - if (run.start < 0 || run.start + run.len > slots()) { - throw std::runtime_error("Pool::read: run out of bounds"); - } - return ::mlx::core::slice( - buf_, - ::mlx::core::Shape{0, 0, run.start, 0}, - ::mlx::core::Shape{1, H, run.start + run.len, D}, - ::mlx::core::Shape{1, 1, 1, 1}, - s); - } - - // Slots currently allocated; grows toward max_slots on demand. - int slots() const { - return static_cast(buf_.shape(2)); - } - - private: - // Make room for `needed` slots, growing only if the pool is short: double - // until it fits, never past max_slots_. Cells keep their index, so growth is - // a zero-pad on the cell axis. - void maybe_grow(int needed, StreamOrDevice s) { - const int cur = slots(); - if (needed <= cur) { - return; - } - int next = std::max(cur, 1); // an empty pool has nothing to double - while (next < needed) { - next *= 2; - } - // The last doubling can overshoot; write() already bounds `needed` by - // max_slots_, so clamping here cannot undershoot it. - next = std::min(next, max_slots_); - const int H = static_cast(buf_.shape(1)); - const int D = static_cast(buf_.shape(3)); - Tensor pad = - ::mlx::core::zeros(::mlx::core::Shape{1, H, next - cur, D}, dtype_); - buf_ = ::mlx::core::concatenate(std::vector{buf_, pad}, 2, s); - } - - ::mlx::core::Dtype dtype_; - int max_slots_; - Tensor buf_; -}; - // Bool mask [1, 1, T, S] for T queries over a span of S keys, where each query // attends at most `window` keys ending at itself. The span is right-aligned // (the newest key belongs to the last query), so query i spans keys @@ -172,7 +81,8 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { AttendSpec update_and_fetch( int layer, - int position, + const int32_t* positions, + int length, const Tensor& k, const Tensor& v, StreamOrDevice s) override { @@ -181,7 +91,8 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { } const int T = static_cast(k.shape(2)); // BHSD: seq axis is 2 - std::optional p = this->plan(layer, position, T); + std::optional p = + this->plan(layer, run_start(positions, length, T), T); if (!p) { throw std::runtime_error( "update_and_fetch: step exceeds capacity or invalid layer"); @@ -216,6 +127,24 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { } private: + // A sequence cache holds one run of one sequence, so the step is described by + // where it starts; the remaining positions carry no information beyond + // confirming that. A step this cache cannot represent is refused rather than + // silently stored at the wrong positions. + static int run_start(const int32_t* positions, int length, int T) { + if (length != T) { + throw std::runtime_error( + "update_and_fetch: one position per query token expected"); + } + for (int i = 1; i < length; ++i) { + if (positions[i] != positions[0] + i) { + throw std::runtime_error( + "update_and_fetch: sequence cache needs a contiguous run"); + } + } + return positions[0]; + } + // Scatter `update` across the step's runs. Runs are in logical order, so // consecutive slices of `update` map to consecutive runs. A flat step is one // run; a ring step splits in two when it wraps the pool. @@ -226,7 +155,7 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { const Tensor& update, StreamOrDevice s) { if (n == 1) { - pool.write(runs[0], update, s); + pool.write(runs[0].start, runs[0].len, update, s); return; } const int H = static_cast(update.shape(1)); @@ -234,7 +163,8 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { int off = 0; for (int i = 0; i < n; ++i) { pool.write( - runs[i], + runs[i].start, + runs[i].len, ::mlx::core::slice( update, ::mlx::core::Shape{0, 0, off, 0}, @@ -250,12 +180,12 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { static Tensor read_runs(const Pool& pool, const cache::Run* runs, int n, StreamOrDevice s) { if (n == 1) { - return pool.read(runs[0], s); + return pool.read(runs[0].start, runs[0].len, s); } std::vector parts; parts.reserve(static_cast(n)); for (int i = 0; i < n; ++i) { - parts.push_back(pool.read(runs[i], s)); + parts.push_back(pool.read(runs[i].start, runs[i].len, s)); } return ::mlx::core::concatenate(parts, 2, s); } diff --git a/backends/mlx/test/mlx_sequence_cache_test.cpp b/backends/mlx/test/mlx_sequence_cache_test.cpp index c4357fd179c..69aa77a080c 100644 --- a/backends/mlx/test/mlx_sequence_cache_test.cpp +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -23,6 +23,7 @@ #include +#include #include #include @@ -32,6 +33,21 @@ using ::mlx::core::array; namespace { +// A step's positions: the contiguous run of k.shape(2) tokens starting at +// `start`, which is what every single-sequence step is. +AttendSpec step( + MLXSequenceCache& c, + int layer, + int start, + const array& k, + const array& v, + ::mlx::core::StreamOrDevice s) { + const int T = static_cast(k.shape(2)); + std::vector positions(static_cast(T)); + std::iota(positions.begin(), positions.end(), start); + return c.update_and_fetch(layer, positions.data(), T, k, v, s); +} + // Max absolute difference within tolerance. Computed in float32: item() // reads sizeof(float) bytes, so calling it on an fp16 scalar misreads the // buffer. @@ -103,7 +119,7 @@ TEST_F(MLXSequenceCacheTest, PrefillIsCausal) { array k0 = randn(T0, float16); array v0 = randn(T0, float16); - AttendSpec spec0 = c.update_and_fetch(0, /*position=*/0, k0, v0, s); + AttendSpec spec0 = step(c, 0, /*position=*/0, k0, v0, s); EXPECT_EQ(spec0.kind, AttendSpec::Mask::Causal); EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); EXPECT_TRUE(allclose(spec0.V, v0, 0.0f)); @@ -122,11 +138,11 @@ TEST_F(MLXSequenceCacheTest, DecodeReadsFullHistory) { const int T0 = 4; array k0 = randn(T0, float16); array v0 = randn(T0, float16); - c.update_and_fetch(0, /*position=*/0, k0, v0, s); // prefill + step(c, 0, /*position=*/0, k0, v0, s); // prefill array k1 = randn(1, float16); array v1 = randn(1, float16); - AttendSpec spec1 = c.update_and_fetch(0, /*position=*/T0, k1, v1, s); + AttendSpec spec1 = step(c, 0, /*position=*/T0, k1, v1, s); EXPECT_EQ(spec1.kind, AttendSpec::Mask::None); EXPECT_TRUE( allclose(spec1.K, concatenate(std::vector{k0, k1}, 2, s), 0.0f)); @@ -144,7 +160,32 @@ TEST_F(MLXSequenceCacheTest, StepPastCapacityThrows) { D, static_cast(ScalarType::Half))); array kx = randn(1, float16); - EXPECT_ANY_THROW(c.update_and_fetch(0, /*position=*/32, kx, kx, s)); + EXPECT_ANY_THROW(step(c, 0, /*position=*/32, kx, kx, s)); +} + +// The step carries one position per query token, and this layout can only hold +// a contiguous run of one sequence. Anything else names cells it cannot +// address, so it is refused rather than stored at the wrong positions. +TEST_F(MLXSequenceCacheTest, NonContiguousOrMiscountedPositionsThrow) { + using namespace ::mlx::core; + MLXSequenceCache c(flat_config( + /*capacity=*/32, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + array k = randn(3, float16); + + const std::vector gap{0, 1, 3}; + EXPECT_ANY_THROW(c.update_and_fetch(0, gap.data(), 3, k, k, s)); + + const std::vector two_seqs{0, 0, 1}; + EXPECT_ANY_THROW(c.update_and_fetch(0, two_seqs.data(), 3, k, k, s)); + + const std::vector short_run{0, 1}; + EXPECT_ANY_THROW(c.update_and_fetch(0, short_run.data(), 2, k, k, s)); + + EXPECT_NO_THROW(step(c, 0, /*position=*/0, k, k, s)); } // Storage dtype != compute: fp32 input, fp16 storage. The cache casts on write, @@ -160,7 +201,7 @@ TEST_F(MLXSequenceCacheTest, StorageDtypeDiffersCastsOnWrite) { const int T0 = 4; array k2 = randn(T0, float32); array v2 = randn(T0, float32); - AttendSpec spec2 = c16.update_and_fetch(0, /*position=*/0, k2, v2, s); + AttendSpec spec2 = step(c16, 0, /*position=*/0, k2, v2, s); EXPECT_EQ(spec2.K.dtype(), float16); EXPECT_EQ(spec2.V.dtype(), float16); EXPECT_TRUE(allclose(spec2.K, astype(k2, float16, s), 0.0f)); @@ -174,12 +215,12 @@ TEST_F(MLXSequenceCacheTest, PoolHonorsRunStart) { using namespace ::mlx::core; Pool p(/*initial_slots=*/8, /*max_slots=*/8, H, D, float16); array x = randn(3, float16); - p.write(cache::Run{/*start=*/2, /*len=*/3}, x, s); + p.write(/*start=*/2, /*len=*/3, x, s); - EXPECT_TRUE(allclose(p.read(cache::Run{2, 3}, s), x, 0.0f)); + EXPECT_TRUE(allclose(p.read(2, 3, s), x, 0.0f)); // The cells before the run are untouched, so reading from 0 is not the same // window -- the regression this guards against. - EXPECT_FALSE(allclose(p.read(cache::Run{0, 3}, s), x, 0.0f)); + EXPECT_FALSE(allclose(p.read(0, 3, s), x, 0.0f)); } // A partial per-layer list is rejected instead of indexing past the end. @@ -209,7 +250,7 @@ TEST_F(MLXSequenceCacheTest, GrowsPastInitialCapacity) { const int T0 = 5; // > initial_capacity array k0 = randn(T0, float16); array v0 = randn(T0, float16); - AttendSpec spec0 = c.update_and_fetch(0, /*position=*/0, k0, v0, s); + AttendSpec spec0 = step(c, 0, /*position=*/0, k0, v0, s); EXPECT_EQ(spec0.K.shape(2), T0); EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); EXPECT_TRUE(allclose(spec0.V, v0, 0.0f)); @@ -229,11 +270,11 @@ TEST_F(MLXSequenceCacheTest, GrowthPreservesExistingCells) { array k0 = randn(2, float16); // exactly fills the initial allocation array v0 = randn(2, float16); - c.update_and_fetch(0, /*position=*/0, k0, v0, s); + step(c, 0, /*position=*/0, k0, v0, s); array k1 = randn(1, float16); // crosses the boundary -> grows array v1 = randn(1, float16); - AttendSpec spec1 = c.update_and_fetch(0, /*position=*/2, k1, v1, s); + AttendSpec spec1 = step(c, 0, /*position=*/2, k1, v1, s); EXPECT_EQ(spec1.K.shape(2), 3); EXPECT_TRUE( allclose(spec1.K, concatenate(std::vector{k0, k1}, 2, s), 0.0f)); @@ -247,12 +288,12 @@ TEST_F(MLXSequenceCacheTest, PoolDoublesAndClampsToMaxSlots) { using namespace ::mlx::core; Pool p(/*initial_slots=*/2, /*max_slots=*/32, H, D, float16); EXPECT_EQ(p.slots(), 2); - p.write(cache::Run{0, 5}, randn(5, float16), s); // 2 -> 4 -> 8 + p.write(0, 5, randn(5, float16), s); // 2 -> 4 -> 8 EXPECT_EQ(p.slots(), 8); // 16 -> 32 overshoots a cap of 20, so it clamps. Pool q(/*initial_slots=*/16, /*max_slots=*/20, H, D, float16); - q.write(cache::Run{0, 17}, randn(17, float16), s); + q.write(0, 17, randn(17, float16), s); EXPECT_EQ(q.slots(), 20); // initial_slots above the cap is clamped at construction. @@ -272,7 +313,7 @@ TEST_F(MLXSequenceCacheTest, ZeroInitialCapacityGrowsOnFirstWrite) { /*initial_capacity=*/0)); array k0 = randn(3, float16); array v0 = randn(3, float16); - AttendSpec spec0 = c.update_and_fetch(0, /*position=*/0, k0, v0, s); + AttendSpec spec0 = step(c, 0, /*position=*/0, k0, v0, s); EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); } @@ -323,7 +364,7 @@ TEST_F(MLXSequenceCacheTest, RingDecodeEvictsOldestAndNeedsNoMask) { } AttendSpec spec{toks[0], toks[0], AttendSpec::Mask::None, {}}; for (int i = 0; i < 6; ++i) { - spec = c.update_and_fetch(0, /*position=*/i, toks[i], toks[i], s); + spec = step(c, 0, /*position=*/i, toks[i], toks[i], s); } EXPECT_EQ(spec.kind, AttendSpec::Mask::None); EXPECT_EQ(spec.K.shape(2), W); @@ -345,7 +386,7 @@ TEST_F(MLXSequenceCacheTest, WindowWiderThanSpanStaysCausal) { D, static_cast(ScalarType::Half))); array k = randn(2, float16); // span 2 < window 4 - AttendSpec rspec = ring.update_and_fetch(0, /*position=*/0, k, k, s); + AttendSpec rspec = step(ring, 0, /*position=*/0, k, k, s); EXPECT_EQ(rspec.kind, AttendSpec::Mask::Causal); EXPECT_FALSE(rspec.mask.has_value()); @@ -355,7 +396,7 @@ TEST_F(MLXSequenceCacheTest, WindowWiderThanSpanStaysCausal) { H, D, static_cast(ScalarType::Half))); - AttendSpec fspec = flat.update_and_fetch(0, /*position=*/0, k, k, s); + AttendSpec fspec = step(flat, 0, /*position=*/0, k, k, s); EXPECT_EQ(fspec.kind, AttendSpec::Mask::Causal); EXPECT_FALSE(fspec.mask.has_value()); } @@ -378,10 +419,10 @@ TEST_F(MLXSequenceCacheTest, RingStepWrapsAndRejoinsInOrder) { std::vector toks; for (int i = 0; i < 4; ++i) { toks.push_back(randn(1, float16)); - c.update_and_fetch(0, /*position=*/i, toks.back(), toks.back(), s); + step(c, 0, /*position=*/i, toks.back(), toks.back(), s); } array pair = randn(2, float16); - AttendSpec spec = c.update_and_fetch(0, /*position=*/4, pair, pair, s); + AttendSpec spec = step(c, 0, /*position=*/4, pair, pair, s); // The span is the union of the two queries' windows -- position 4 attends // 1..4 and position 5 attends 2..5 -- so it is window + T - 1 = 5 cells, not