[MLX] Per-token positions and a shared pool

backends/mlx/runtime/MLXCache.h

@@ -8,7 +8,9 @@
 
 #pragma once
 
+#include <cstdint>
 #include <optional>
+#include <vector>
 
 #include "MLXExecutor.h" // Tensor, StreamOrDevice
 
@@ -36,14 +38,11 @@ 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 std::vector<int32_t>& positions,
       const Tensor& k,
       const Tensor& v,
       StreamOrDevice s) = 0;

backends/mlx/runtime/MLXInterpreter.h

@@ -11,6 +11,9 @@
 #include "MLXCache.h"
 #include "MLXExecutor.h"
 
+#include <algorithm>
+#include <vector>
+
 #include <mlx/array.h>
 #include <mlx/fast.h>
 #include <mlx/mlx.h>
@@ -310,20 +313,34 @@ 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.
+  // nothing for the rest.
   auto pos = st.const_tensor_ref(n.position);
   eval(pos);
-  int position;
+  // The entries are read in order off the buffer, which a strided view would
+  // walk with the wrong stride.
+  if (!pos.flags().row_contiguous) {
+    throw std::runtime_error("update_and_attend: position must be contiguous");
+  }
+  const int length = static_cast<int>(pos.size());
+  if (length != static_cast<int>(q.shape(2))) {
+    throw std::runtime_error(
+        "update_and_attend: position must hold one entry per query token");
+  }
+  std::vector<int32_t> positions(static_cast<size_t>(length));
   switch (pos.dtype()) {
     case ::mlx::core::int32:
-      position = pos.data<int32_t>()[0];
+      std::copy(
+          pos.data<int32_t>(), pos.data<int32_t>() + length, positions.begin());
       break;
     case ::mlx::core::int64:
-      position = static_cast<int>(pos.data<int64_t>()[0]);
+      std::transform(
+          pos.data<int64_t>(),
+          pos.data<int64_t>() + length,
+          positions.begin(),
+          [](int64_t p) { return static_cast<int32_t>(p); });
       break;
     default:
       throw std::runtime_error(
@@ -332,18 +349,17 @@ inline void exec_update_and_attend(
   }
   AttendSpec spec = st.cache->update_and_fetch(
       *n.layer_id,
-      position,
+      positions,
       st.const_tensor_ref(n.k),
       st.const_tensor_ref(n.v),
       s);
   // Match stored K/V to the query dtype before SDPA (no-op when equal; the
   // storage precision may differ from the compute dtype).
   array K = spec.K.dtype() == q.dtype() ? spec.K : astype(spec.K, q.dtype(), s);
   array V = spec.V.dtype() == q.dtype() ? spec.V : astype(spec.V, q.dtype(), s);
-  // MLX takes the mask as a mode string plus an optional tensor. Switch rather
-  // than test for Causal: None and Explicit both map to "" and are told apart
-  // only by spec.mask, so an Explicit with no mask would silently attend
-  // unmasked.
+  // MLX takes the mask as a mode string plus an optional tensor. Switch: None
+  // and Explicit both map to "" and are told apart only by spec.mask, so an
+  // Explicit with no mask would silently attend unmasked.
   std::string mask_mode;
   switch (spec.kind) {
     case AttendSpec::Mask::None:

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 <algorithm>
+#include <stdexcept>
+#include <vector>
+
+#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<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) {
+      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<int>(buf_.shape(1));
+    const int D = static_cast<int>(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<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_;
+};
+
+} // namespace mlx
+} // namespace backends
+} // namespace executorch