diff --git a/backends/mlx/runtime/MLXCache.h b/backends/mlx/runtime/MLXCache.h index 8c4e45d8cbe..98f144cfdf9 100644 --- a/backends/mlx/runtime/MLXCache.h +++ b/backends/mlx/runtime/MLXCache.h @@ -16,11 +16,10 @@ namespace executorch { namespace backends { namespace mlx { -// The K/V window to attend over + how to mask it; the cache owns the semantic. -// `kind` mirrors MLX SDPA's mask forms (no mask / "causal" / explicit tensor). -// Will be added: a `window` on Causal (sliding-window -- the -// mask_mod axis) and a `softcap`/bias field (ALiBi, softcap -- the score_mod -// axis). +// The K/V window to attend over + how to mask it; the cache owns the semantic +// and hands over whatever MLX SDPA needs to apply it. `kind` mirrors MLX's mask +// forms: no mask, its fused "causal", or an explicit tensor for anything MLX +// cannot express -- a sliding window, and later tree/speculative patterns. struct AttendSpec { Tensor K; Tensor V; diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 3ca4041147a..9e644c3d977 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -117,15 +117,33 @@ class Pool { 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 +// j - i <= S - T -- the same bound MLX's "causal" applies -- and the window +// adds the lower bound j - i > S - T - window. +inline Tensor window_causal_mask(int T, int S, int window, StreamOrDevice s) { + using namespace ::mlx::core; + Tensor diff = subtract( + reshape(arange(S, int32, s), Shape{1, S}, s), + reshape(arange(T, int32, s), Shape{T, 1}, s), + s); + const int hi = S - T; + Tensor band = logical_and( + less_equal(diff, array(hi), s), + greater_equal(diff, array(hi - window + 1), s), + s); + return reshape(band, Shape{1, 1, T, S}, s); +} + // The MLX byte layer behind the neutral SequenceCache: one Pool per layer for K // and one for V, sized from that layer's own policy. update_and_fetch plans the // step (integer runs), writes the new K/V, reads the retained window, and // declares the mask; the op handler owns q/scale and calls SDPA. // -// Only flat layers are wired up so far. Ring layers are rejected at -// construction rather than silently mis-served: the plumbing below handles -// their runs, but the windowed mask and read_base_pos (RoPE alignment) are not -// implemented yet. +// Flat and ring layers can be mixed (gemma4 alternates them): the planner emits +// physical runs either way, so they differ here only in how many slots the pool +// holds and whether a step's runs wrap. class MLXSequenceCache : public cache::SequenceCache, public MLXCache { public: explicit MLXSequenceCache(const cache::CacheConfig& cfg) @@ -134,21 +152,21 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { resolve_dtype(static_cast(cfg.kv_dtype)); kpool_.reserve(static_cast(cfg.n_layers)); vpool_.reserve(static_cast(cfg.n_layers)); + window_.reserve(static_cast(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 cache::LayerConfig& lc = cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; - if (lc.policy.kind != cache::LayerPolicy::Kind::Flat) { - throw std::runtime_error( - "MLXSequenceCache: only Flat layers are supported so far"); - } - // Flat retains all history, so the layer's pool may reach the full cap. A - // ring layer will instead cap at window + max_write - 1 slots. - const int max_slots = cfg.capacity; - kpool_.emplace_back( - cfg.initial_capacity, max_slots, lc.n_kv_heads, lc.head_dim, dt); - vpool_.emplace_back( - cfg.initial_capacity, max_slots, lc.n_kv_heads, lc.head_dim, dt); + const bool ring = lc.policy.kind == cache::LayerPolicy::Kind::Ring; + window_.push_back(ring ? lc.policy.window : 0); + // Flat retains all history, so its pool may reach the full cap and starts + // small. A ring layer recycles a fixed window + max_write - 1 slots. + const int max_slots = ring ? lc.policy.window + + (cfg.max_write ? *cfg.max_write : lc.policy.window) - 1 + : cfg.capacity; + const int initial = ring ? max_slots : cfg.initial_capacity; + kpool_.emplace_back(initial, max_slots, lc.n_kv_heads, lc.head_dim, dt); + vpool_.emplace_back(initial, max_slots, lc.n_kv_heads, lc.head_dim, dt); } } @@ -168,27 +186,80 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { throw std::runtime_error( "update_and_fetch: step exceeds capacity or invalid layer"); } - // Flat layers never wrap, and ring layers are refused at construction, so - // every accepted step is a single run. Ring's split/rejoin lands with ring. - if (p->n_write != 1 || p->n_read != 1) { - throw std::runtime_error("update_and_fetch: expected a single-run plan"); - } const size_t l = static_cast(layer); - kpool_[l].write(p->write[0], k, s); - vpool_[l].write(p->write[0], v, s); - Tensor K = kpool_[l].read(p->read[0], s); - Tensor V = vpool_[l].read(p->read[0], s); + write_runs(kpool_[l], p->write, p->n_write, k, s); + write_runs(vpool_[l], p->write, p->n_write, v, s); + Tensor K = read_runs(kpool_[l], p->read, p->n_read, s); + Tensor V = read_runs(vpool_[l], p->read, p->n_read, s); this->commit(*p); - // A multi-token chain is Causal (MLX "causal" is lower-right aligned, so - // fresh and chunked prefill are both correct with new tokens at the tail); - // a single decode token needs no mask. - const AttendSpec::Mask kind = - (T > 1) ? AttendSpec::Mask::Causal : AttendSpec::Mask::None; - return AttendSpec{K, V, kind, std::nullopt}; + // A single decode token needs no mask -- its span is exactly what it may + // attend, on a ring layer as much as a flat one. + if (T == 1) { + return AttendSpec{K, V, AttendSpec::Mask::None, std::nullopt}; + } + // A ring layer bounds each query to its own window. That only bites when + // the window is narrower than the span (a multi-token step reads the union + // of its queries' windows, window + T - 1 cells); otherwise plain causal is + // exact and MLX applies it fused, with no mask tensor. + const int S = static_cast(K.shape(2)); + if (window_[l] > 0 && window_[l] < S) { + return AttendSpec{ + K, + V, + AttendSpec::Mask::Explicit, + window_causal_mask(T, S, window_[l], s)}; + } + // MLX "causal" is lower-right aligned, so fresh and chunked prefill are + // both correct with the new tokens at the tail. + return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt}; } private: + // 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. + static void write_runs( + Pool& pool, + const cache::Run* runs, + int n, + const Tensor& update, + StreamOrDevice s) { + if (n == 1) { + pool.write(runs[0], update, s); + return; + } + const int H = static_cast(update.shape(1)); + const int D = static_cast(update.shape(3)); + int off = 0; + for (int i = 0; i < n; ++i) { + pool.write( + runs[i], + ::mlx::core::slice( + update, + ::mlx::core::Shape{0, 0, off, 0}, + ::mlx::core::Shape{1, H, off + runs[i].len, D}, + ::mlx::core::Shape{1, 1, 1, 1}, + s), + s); + off += runs[i].len; + } + } + + // Gather the step's runs into the retained window, oldest -> newest. + static Tensor + read_runs(const Pool& pool, const cache::Run* runs, int n, StreamOrDevice s) { + if (n == 1) { + return pool.read(runs[0], s); + } + std::vector parts; + parts.reserve(static_cast(n)); + for (int i = 0; i < n; ++i) { + parts.push_back(pool.read(runs[i], s)); + } + return ::mlx::core::concatenate(parts, 2, s); + } + // Enforce the neutral contract as an exception, the failure mode this layer // already uses. Runs as the base initializer's argument because // SequenceCache's own ctor indexes `layers` before this class's body does. @@ -201,6 +272,7 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { std::vector kpool_; std::vector vpool_; + std::vector window_; // per layer; 0 = flat (unbounded history) }; } // namespace mlx diff --git a/backends/mlx/test/mlx_sequence_cache_test.cpp b/backends/mlx/test/mlx_sequence_cache_test.cpp index 0dff49acfcc..c4357fd179c 100644 --- a/backends/mlx/test/mlx_sequence_cache_test.cpp +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -63,6 +63,22 @@ cache::CacheConfig flat_config( return cfg; } +// One ring layer of `window` cells; max_write bounds a multi-token step. +cache::CacheConfig ring_config( + int capacity, + int window, + int max_write, + int n_kv_heads, + int head_dim, + int kv_dtype) { + cache::CacheConfig cfg = + flat_config(capacity, /*n_layers=*/1, n_kv_heads, head_dim, kv_dtype); + cfg.layers[0].policy = + cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, window}; + cfg.max_write = max_write; + return cfg; +} + class MLXSequenceCacheTest : public ::testing::Test { protected: const int H = 2; @@ -260,6 +276,137 @@ TEST_F(MLXSequenceCacheTest, ZeroInitialCapacityGrowsOnFirstWrite) { EXPECT_TRUE(allclose(spec0.K, k0, 0.0f)); } +// The sliding-window mask is a band on (key - query): causal above, window +// below. The span is right-aligned, so the last query's own key is the last +// key, and each query attends `window` keys ending at its own. +TEST_F(MLXSequenceCacheTest, WindowCausalMaskIsABand) { + using namespace ::mlx::core; + // T=3 queries over S=5 keys, window 3. Query i owns key i + (S - T), and + // attends the `window` keys ending there -- a band, one row per query. + // clang-format off + std::vector want = { + 1, 1, 1, 0, 0, + 0, 1, 1, 1, 0, + 0, 0, 1, 1, 1}; + // A window covering the whole span leaves plain causal: no lower bound bites. + std::vector causal = { + 1, 1, 1, 0, 0, + 1, 1, 1, 1, 0, + 1, 1, 1, 1, 1}; + // clang-format on + array m = astype(window_causal_mask(3, 5, 3, s), int32, s); + EXPECT_TRUE(allclose(m, array(want.data(), Shape{1, 1, 3, 5}, int32), 0.0f)); + + array full = astype(window_causal_mask(3, 5, 5, s), int32, s); + EXPECT_TRUE( + allclose(full, array(causal.data(), Shape{1, 1, 3, 5}, int32), 0.0f)); +} + +// A ring layer decodes past its window: the read span stops at `window` cells +// and holds the newest tokens, evicting the oldest. No mask -- a single query +// may attend its whole span. +TEST_F(MLXSequenceCacheTest, RingDecodeEvictsOldestAndNeedsNoMask) { + using namespace ::mlx::core; + const int W = 4; + MLXSequenceCache c(ring_config( + /*capacity=*/64, + /*window=*/W, + /*max_write=*/1, + H, + D, + static_cast(ScalarType::Half))); + + // Feed 6 single tokens; the last 4 must survive, oldest -> newest. + std::vector toks; + for (int i = 0; i < 6; ++i) { + toks.push_back(randn(1, float16)); + } + 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); + } + EXPECT_EQ(spec.kind, AttendSpec::Mask::None); + EXPECT_EQ(spec.K.shape(2), W); + EXPECT_TRUE(allclose( + spec.K, + concatenate(std::vector{toks[2], toks[3], toks[4], toks[5]}, 2, s), + 0.0f)); +} + +// A ring layer stays Causal while its window still covers the span, so MLX +// applies its fused mask and no tensor is built. A flat layer always does. +TEST_F(MLXSequenceCacheTest, WindowWiderThanSpanStaysCausal) { + using namespace ::mlx::core; + MLXSequenceCache ring(ring_config( + /*capacity=*/64, + /*window=*/4, + /*max_write=*/2, + H, + 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); + EXPECT_EQ(rspec.kind, AttendSpec::Mask::Causal); + EXPECT_FALSE(rspec.mask.has_value()); + + MLXSequenceCache flat(flat_config( + /*capacity=*/64, + /*n_layers=*/1, + H, + D, + static_cast(ScalarType::Half))); + AttendSpec fspec = flat.update_and_fetch(0, /*position=*/0, k, k, s); + EXPECT_EQ(fspec.kind, AttendSpec::Mask::Causal); + EXPECT_FALSE(fspec.mask.has_value()); +} + +// A step whose runs wrap the ring is scattered and gathered in logical order. +TEST_F(MLXSequenceCacheTest, RingStepWrapsAndRejoinsInOrder) { + using namespace ::mlx::core; + const int W = 4; + const int MW = 2; + MLXSequenceCache c(ring_config( + /*capacity=*/64, + /*window=*/W, + /*max_write=*/MW, + H, + D, + static_cast(ScalarType::Half))); + + // ring_size = W + MW - 1 = 5. Fill 4 tokens, then a 2-token step at + // position 4 writes slots 4 and 0 -- a wrap. + 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); + } + array pair = randn(2, float16); + AttendSpec spec = c.update_and_fetch(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 + // 4. Since the window is now narrower than the span, the cache hands back a + // band to narrow each query back to its own 4. + EXPECT_EQ(spec.K.shape(2), W + 1); + EXPECT_EQ(spec.kind, AttendSpec::Mask::Explicit); + ASSERT_TRUE(spec.mask.has_value()); + // Span indices 0..4 are positions 1..5. Query 0 is position 4 and attends + // 1..4; query 1 is position 5 and attends 2..5 -- four keys each. + // clang-format off + std::vector band = { + 1, 1, 1, 1, 0, + 0, 1, 1, 1, 1}; + // clang-format on + EXPECT_TRUE(allclose( + astype(*spec.mask, int32, s), + array(band.data(), Shape{1, 1, 2, W + 1}, int32), + 0.0f)); + EXPECT_TRUE(allclose( + spec.K, + concatenate(std::vector{toks[1], toks[2], toks[3], pair}, 2, s), + 0.0f)); +} + // A negative initial_capacity is rejected rather than reaching MLX as a // negative dimension. TEST_F(MLXSequenceCacheTest, NegativeInitialCapacityThrows) {