diff --git a/backends/mlx/runtime/MLXCache.h b/backends/mlx/runtime/MLXCache.h index 8c4e45d8cbe..55440b081b2 100644 --- a/backends/mlx/runtime/MLXCache.h +++ b/backends/mlx/runtime/MLXCache.h @@ -17,15 +17,19 @@ 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). +// `kind` mirrors MLX SDPA's mask forms (no mask / "causal" / explicit tensor), +// with `window` as the declarative sliding-window modifier (the mask_mod axis). +// The cache states the semantic; the handler materializes whatever the backend +// cannot express natively. Will be added: a `softcap`/bias field (ALiBi, +// softcap -- the score_mod axis). struct AttendSpec { Tensor K; Tensor V; enum class Mask { None, Causal, Explicit } kind; std::optional mask; // Explicit only + // Causal only: each query attends at most this many keys (itself included). + // Unset = unbounded history. + std::optional window; }; // Tensor-typed op face of the off-graph KV cache, kept separate from the diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 51946e966a6..f785dc41bcf 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -294,6 +294,24 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor(n.out, std::move(out)); } +// 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 array window_causal_mask(int T, int S, int window, StreamOrDevice s) { + array 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; + array 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); +} + inline void exec_update_and_attend( const UpdateAndAttendNode& n, ExecutionState& st, @@ -333,12 +351,22 @@ inline void exec_update_and_attend( // only by spec.mask, so an Explicit with no mask would silently attend // unmasked. std::string mask_mode; + std::optional mask = spec.mask; switch (spec.kind) { case AttendSpec::Mask::None: break; - case AttendSpec::Mask::Causal: - mask_mode = "causal"; + case AttendSpec::Mask::Causal: { + const int T = static_cast(q.shape(2)); + const int S = static_cast(K.shape(2)); + // MLX has no sliding-window mode, so a window that actually bites has to + // be materialized. It never bites when it covers the whole span. + if (spec.window && *spec.window < S) { + mask = window_causal_mask(T, S, *spec.window, s); + } else { + mask_mode = "causal"; + } break; + } case AttendSpec::Mask::Explicit: if (!spec.mask) { throw std::runtime_error( @@ -347,14 +375,7 @@ inline void exec_update_and_attend( break; } array out = fast::scaled_dot_product_attention( - q, - K, - V, - static_cast(*n.scale), - mask_mode, - spec.mask, - std::nullopt, - s); + q, K, V, static_cast(*n.scale), mask_mode, mask, std::nullopt, s); // Honor the op's output-dtype contract (unset -> SDPA's native output). if (n.out_dtype) { out = astype(out, resolve_dtype(*n.out_dtype), s); diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 3ca4041147a..a8695631e9c 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -122,10 +122,9 @@ class Pool { // 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 +133,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 +167,75 @@ 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. A ring layer additionally + // bounds each query to its own window; the handler materializes that only + // when the window is narrower than the span. + if (T == 1) { + return AttendSpec{ + K, V, AttendSpec::Mask::None, std::nullopt, std::nullopt}; + } + std::optional window; + if (window_[l] > 0) { + window = window_[l]; + } + return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt, window}; } 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 +248,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..372c460c615 100644 --- a/backends/mlx/test/mlx_sequence_cache_test.cpp +++ b/backends/mlx/test/mlx_sequence_cache_test.cpp @@ -17,6 +17,7 @@ // // Must run on Apple Silicon: MLX needs the Metal backend. +#include "MLXInterpreter.h" // window_causal_mask #include "MLXSequenceCache.h" #include @@ -63,6 +64,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 +277,127 @@ 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(ops::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(ops::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's multi-token step declares its window; the handler decides +// whether to materialize it. Flat layers declare no window. +TEST_F(MLXSequenceCacheTest, RingDeclaresWindowOnCausal) { + using namespace ::mlx::core; + const int W = 4; + MLXSequenceCache ring(ring_config( + /*capacity=*/64, + /*window=*/W, + /*max_write=*/2, + H, + D, + static_cast(ScalarType::Half))); + array k = randn(2, float16); + AttendSpec rspec = ring.update_and_fetch(0, /*position=*/0, k, k, s); + EXPECT_EQ(rspec.kind, AttendSpec::Mask::Causal); + ASSERT_TRUE(rspec.window.has_value()); + EXPECT_EQ(*rspec.window, W); + + 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.window.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. Narrowing each query back to its own 4 is the declared window's job. + EXPECT_EQ(spec.K.shape(2), W + 1); + ASSERT_TRUE(spec.window.has_value()); + EXPECT_EQ(*spec.window, W); + 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) {