[MLX] Place the KV pool's arrays on the stream it was given (pytorch#22230)

Pool's constructor allocated buf_ with a zeros() that took no stream, so
it fell back to default_stream(): the calling thread's default, resolved
at call time and stamped into the array's primitive.

The constructor cannot be fixed by passing it a stream, because none
exists yet. The stream belongs to the delegate handle, and the cache is
built and registered under a key before init() constructs the handle
that looks it up.
So the constructor now allocates nothing: Pool holds its geometry and
materialises buf_ on the first maybe_grow, on the stream that write
supplied. That leaves one place in the class that creates arrays, and it
cannot run without a stream.

slots() now reports 0 until the first write, and read() on an untouched
pool throws instead of returning zeros. PoolDoublesAndClampsToMaxSlots
asserts the initial and clamped sizes after a write rather than after
construction.

Author: kiymetakdemir

backends/mlx/runtime/MLXPool.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <algorithm>
+#include <optional>
 #include <stdexcept>
 #include <vector>
 
@@ -25,38 +26,38 @@ namespace mlx {
 // where in them a step lands.
 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.
+  // Holds its geometry but allocates nothing: every array belongs to the
+  // stream its work runs on, and the first stream reaches the pool with the
+  // first write. initial_slots above max_slots is clamped, not rejected: the
+  // config default exceeds the cap of any smaller cache.
   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)) {}
+        initial_slots_(std::min(initial_slots, max_slots)),
+        H_(H),
+        D_(D) {}
 
   // 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<int>(buf_.shape(1));
-    const int D = static_cast<int>(buf_.shape(3));
     if (start < 0 || start + len > max_slots_) {
       throw std::runtime_error("Pool::write: run out of bounds");
     }
     if (static_cast<int>(update.shape(2)) != len) {
       throw std::runtime_error("Pool::write: update length != run length");
     }
-    if (static_cast<int>(update.shape(1)) != H ||
-        static_cast<int>(update.shape(3)) != D) {
+    if (static_cast<int>(update.shape(1)) != H_ ||
+        static_cast<int>(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_,
+        *buf_,
         u,
         ::mlx::core::Shape{0, 0, start, 0},
-        ::mlx::core::Shape{1, H, start + len, D},
+        ::mlx::core::Shape{1, H_, start + len, D_},
         s);
   }
 
@@ -67,15 +68,13 @@ class Pool {
       const std::vector<int32_t>& cells,
       const Tensor& update,
       StreamOrDevice s) {
-    const int H = static_cast<int>(buf_.shape(1));
-    const int D = static_cast<int>(buf_.shape(3));
     const int T = static_cast<int>(cells.size());
     if (static_cast<int>(update.shape(2)) != T) {
       throw std::runtime_error(
           "Pool::write_cells: update length != cell count");
     }
-    if (static_cast<int>(update.shape(1)) != H ||
-        static_cast<int>(update.shape(3)) != D) {
+    if (static_cast<int>(update.shape(1)) != H_ ||
+        static_cast<int>(update.shape(3)) != D_) {
       throw std::runtime_error("Pool::write_cells: K/V heads/dim mismatch");
     }
     int high = 0;
@@ -92,7 +91,7 @@ class Pool {
     // put_along_axis broadcasts the one index per token across heads and
     // head_dim.
     buf_ = ::mlx::core::put_along_axis(
-        buf_,
+        *buf_,
         ::mlx::core::array(
             cells.data(), ::mlx::core::Shape{1, 1, T, 1}, ::mlx::core::int32),
         u,
@@ -103,29 +102,33 @@ class Pool {
   // 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<int>(buf_.shape(1));
-    const int D = static_cast<int>(buf_.shape(3));
-    if (start < 0 || start + len > slots()) {
+    if (!buf_ || start < 0 || start + len > slots()) {
       throw std::runtime_error("Pool::read: run out of bounds");
     }
     return ::mlx::core::slice(
-        buf_,
+        *buf_,
         ::mlx::core::Shape{0, 0, start, 0},
-        ::mlx::core::Shape{1, H, start + len, D},
+        ::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.
+  // Slots currently allocated; 0 until the first write, then grows toward
+  // max_slots on demand.
   int slots() const {
-    return static_cast<int>(buf_.shape(2));
+    return buf_ ? static_cast<int>(buf_->shape(2)) : 0;
   }
 
  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.
+  // Make room for `needed` slots, allocating on first use and thereafter
+  // 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. Every array here is placed on `s`, the stream the step runs on.
   void maybe_grow(int needed, StreamOrDevice s) {
+    if (!buf_) {
+      buf_ = ::mlx::core::zeros(
+          ::mlx::core::Shape{1, H_, initial_slots_, D_}, dtype_, s);
+    }
     const int cur = slots();
     if (needed <= cur) {
       return;
@@ -137,16 +140,18 @@ class Pool {
     // 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);
+    Tensor pad = ::mlx::core::zeros(
+        ::mlx::core::Shape{1, H_, next - cur, D_}, dtype_, s);
+    buf_ = ::mlx::core::concatenate(std::vector<Tensor>{*buf_, pad}, 2, s);
   }
 
   ::mlx::core::Dtype dtype_;
   int max_slots_;
-  Tensor buf_;
+  int initial_slots_;
+  int H_;
+  int D_;
+  // Absent until the first write hands the pool a stream to allocate on.
+  std::optional<Tensor> buf_;
 };
 
 } // namespace mlx

backends/mlx/test/mlx_sequence_cache_test.cpp

@@ -241,7 +241,7 @@ TEST_F(MLXSequenceCacheTest, GrowthPreservesExistingCells) {
 TEST_F(MLXSequenceCacheTest, PoolDoublesAndClampsToMaxSlots) {
   using namespace ::mlx::core;
   Pool p(/*initial_slots=*/2, /*max_slots=*/32, H, D, float16);
-  EXPECT_EQ(p.slots(), 2);
+  EXPECT_EQ(p.slots(), 0) << "a pool holds no slots until its first write";
   p.write(0, 5, randn(5, float16), s); // 2 -> 4 -> 8
   EXPECT_EQ(p.slots(), 8);
 
@@ -250,8 +250,9 @@ TEST_F(MLXSequenceCacheTest, PoolDoublesAndClampsToMaxSlots) {
   q.write(0, 17, randn(17, float16), s);
   EXPECT_EQ(q.slots(), 20);
 
-  // initial_slots above the cap is clamped at construction.
+  // initial_slots above the cap is clamped when the pool allocates.
   Pool r(/*initial_slots=*/512, /*max_slots=*/4, H, D, float16);
+  r.write(0, 1, randn(1, float16), s);
   EXPECT_EQ(r.slots(), 4);
 }