[MLX] Grow the KV-cache pool on demand (#21525)

Summary

Makes the MLX KV-cache pool grow on demand instead of allocating the
full cap at construction. A pool now starts at
CacheConfig::initial_capacity and doubles as writes need room, up to the
layer's cap; so a short session doesn't pay for a long one's worth of
memory.

Files

- backends/mlx/runtime/MLXSequenceCache.h — Pool takes both bounds
(initial_slots, max_slots) and allocates min of
them. A write bounds its run against the cap, then grow_for doubles from
the current size and clamps to the cap,
since the last doubling can overshoot. MLXSequenceCache passes
cfg.initial_capacity alongside the per-layer cap.
- extension/llm/cache/cache.h — valid(cfg) now rejects a negative
initial_capacity; zero is allowed and means
  "allocate nothing up front."
  - backends/mlx/test/mlx_sequence_cache_test.cpp — four growth cases.

  Testing

  Four cases added to the existing GTest suite (11 total), all passing:

- a first write larger than the initial allocation grows instead of
failing
- a decode crossing the allocated boundary still reads back the full
history (growth preserves written cells)
- doubling stops once the run fits, a 16→32 doubling clamps to a cap of
20, and an initial above the cap clamps at
  construction
- a zero initial capacity grows on first write; a negative one is
rejected

  cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON
  cmake --build cmake-out --target mlx_sequence_cache_test
  ctest --test-dir cmake-out -R mlx_sequence_cache --output-on-failure

Author: kiymetakdemir

backends/mlx/runtime/MLXSequenceCache.h

@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <optional>
 #include <stdexcept>
 #include <string>
@@ -31,16 +32,21 @@ namespace cache = ::executorch::extension::llm::cache;
 // layer asks for and how many runs a step produces.
 class Pool {
  public:
-  Pool(int slots, int H, int D, ::mlx::core::Dtype dtype)
+  // 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),
-        buf_(::mlx::core::zeros(::mlx::core::Shape{1, H, slots, D}, 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<int>(buf_.shape(1));
     const int D = static_cast<int>(buf_.shape(3));
-    if (run.start < 0 || run.start + run.len > slots()) {
+    if (run.start < 0 || run.start + run.len > max_slots_) {
       throw std::runtime_error("Pool::write: run out of bounds");
     }
     if (static_cast<int>(update.shape(2)) != run.len) {
@@ -50,6 +56,7 @@ class Pool {
         static_cast<int>(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);
@@ -77,12 +84,36 @@ class Pool {
         s);
   }
 
+  // Slots currently allocated; grows toward max_slots on demand.
   int slots() const {
     return static_cast<int>(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<int>(buf_.shape(1));
+    const int D = static_cast<int>(buf_.shape(3));
+    Tensor pad =
+        ::mlx::core::zeros(::mlx::core::Shape{1, H, next - cur, D}, dtype_);
+    buf_ = ::mlx::core::concatenate(std::vector<Tensor>{buf_, pad}, 2, s);
+  }
+
   ::mlx::core::Dtype dtype_;
+  int max_slots_;
   Tensor buf_;
 };
 
@@ -111,11 +142,13 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
         throw std::runtime_error(
             "MLXSequenceCache: only Flat layers are supported so far");
       }
-      // Flat retains all history, so the layer's pool is the full cap. A ring
-      // layer will instead ask for window + max_write - 1 slots.
-      const int slots = cfg.capacity;
-      kpool_.emplace_back(slots, lc.n_kv_heads, lc.head_dim, dt);
-      vpool_.emplace_back(slots, lc.n_kv_heads, lc.head_dim, dt);
+      // 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);
     }
   }
 

backends/mlx/test/mlx_sequence_cache_test.cpp

@@ -23,6 +23,7 @@
 
 #include <gtest/gtest.h>
 
+#include <optional>
 #include <vector>
 
 using namespace ::executorch::backends::mlx;
@@ -46,7 +47,8 @@ cache::CacheConfig flat_config(
     int n_layers,
     int n_kv_heads,
     int head_dim,
-    int kv_dtype) {
+    int kv_dtype,
+    std::optional<int> initial_capacity = std::nullopt) {
   cache::CacheConfig cfg;
   cfg.capacity = capacity;
   cfg.n_layers = n_layers;
@@ -55,6 +57,9 @@ cache::CacheConfig flat_config(
       n_kv_heads,
       head_dim}};
   cfg.kv_dtype = kv_dtype;
+  if (initial_capacity) {
+    cfg.initial_capacity = *initial_capacity;
+  }
   return cfg;
 }
 
@@ -151,7 +156,7 @@ TEST_F(MLXSequenceCacheTest, StorageDtypeDiffersCastsOnWrite) {
 // mid-pool, and dropping the start would silently return the wrong cells.
 TEST_F(MLXSequenceCacheTest, PoolHonorsRunStart) {
   using namespace ::mlx::core;
-  Pool p(/*slots=*/8, H, D, float16);
+  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);
 
@@ -173,4 +178,99 @@ TEST_F(MLXSequenceCacheTest, PartialLayerListThrows) {
   EXPECT_ANY_THROW(MLXSequenceCache{cfg});
 }
 
+// A step past the allocated slots grows the pool instead of failing, and the
+// result is the same window a fully-allocated pool would have returned.
+TEST_F(MLXSequenceCacheTest, GrowsPastInitialCapacity) {
+  using namespace ::mlx::core;
+  MLXSequenceCache c(flat_config(
+      /*capacity=*/32,
+      /*n_layers=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half),
+      /*initial_capacity=*/2));
+
+  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);
+  EXPECT_EQ(spec0.K.shape(2), T0);
+  EXPECT_TRUE(allclose(spec0.K, k0, 0.0f));
+  EXPECT_TRUE(allclose(spec0.V, v0, 0.0f));
+}
+
+// Growth preserves cells already written: a decode crossing the allocated
+// boundary must still read back the full history.
+TEST_F(MLXSequenceCacheTest, GrowthPreservesExistingCells) {
+  using namespace ::mlx::core;
+  MLXSequenceCache c(flat_config(
+      /*capacity=*/32,
+      /*n_layers=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half),
+      /*initial_capacity=*/2));
+
+  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);
+
+  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);
+  EXPECT_EQ(spec1.K.shape(2), 3);
+  EXPECT_TRUE(
+      allclose(spec1.K, concatenate(std::vector<array>{k0, k1}, 2, s), 0.0f));
+  EXPECT_TRUE(
+      allclose(spec1.V, concatenate(std::vector<array>{v0, v1}, 2, s), 0.0f));
+}
+
+// Growth doubles until the run fits, and never allocates past max_slots --
+// including when the last doubling would overshoot it.
+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
+  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);
+  EXPECT_EQ(q.slots(), 20);
+
+  // initial_slots above the cap is clamped at construction.
+  Pool r(/*initial_slots=*/512, /*max_slots=*/4, H, D, float16);
+  EXPECT_EQ(r.slots(), 4);
+}
+
+// A pool that starts empty is allowed, and grows on the first write.
+TEST_F(MLXSequenceCacheTest, ZeroInitialCapacityGrowsOnFirstWrite) {
+  using namespace ::mlx::core;
+  MLXSequenceCache c(flat_config(
+      /*capacity=*/32,
+      /*n_layers=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half),
+      /*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);
+  EXPECT_TRUE(allclose(spec0.K, k0, 0.0f));
+}
+
+// A negative initial_capacity is rejected rather than reaching MLX as a
+// negative dimension.
+TEST_F(MLXSequenceCacheTest, NegativeInitialCapacityThrows) {
+  cache::CacheConfig cfg = flat_config(
+      /*capacity=*/32,
+      /*n_layers=*/1,
+      H,
+      D,
+      static_cast<int>(ScalarType::Half),
+      /*initial_capacity=*/-1);
+  EXPECT_ANY_THROW(MLXSequenceCache{cfg});
+}
+
 } // namespace

extension/llm/cache/cache.h

@@ -136,7 +136,9 @@ struct CacheConfig {
 // list that is neither size 1 nor n_layers reads past the end. Reported as a
 // bool rather than thrown, so each backend picks its own failure mode.
 inline bool valid(const CacheConfig& cfg) {
-  return cfg.capacity > 0 && cfg.n_layers > 0 &&
+  // initial_capacity may be 0 (allocate nothing up front) but not negative, and
+  // may exceed capacity -- the byte layer clamps it.
+  return cfg.capacity > 0 && cfg.n_layers > 0 && cfg.initial_capacity >= 0 &&
       (cfg.layers.size() == 1 ||
        cfg.layers.size() == static_cast<size_t>(cfg.n_layers));
 }