[MLX] Per-token positions and a shared pool

backends/mlx/runtime/MLXCache.h

@@ -36,14 +36,12 @@ 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 int32_t* positions,
+      int length,
       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,33 @@ 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.
   auto pos = st.const_tensor_ref(n.position);
   eval(pos);
-  int position;
+  const int length = static_cast<int>(pos.size());
+  if (length != static_cast<int>(st.const_tensor_ref(n.k).shape(2))) {
+    throw std::runtime_error(
+        "update_and_attend: position must hold one entry per query token");
+  }
+  // int32 is passed straight through; only an int64 input needs narrowing.
+  std::vector<int32_t> int32_positions;
+  const int32_t* positions = nullptr;
   switch (pos.dtype()) {
     case ::mlx::core::int32:
-      position = pos.data<int32_t>()[0];
+      positions = pos.data<int32_t>();
       break;
     case ::mlx::core::int64:
-      position = static_cast<int>(pos.data<int64_t>()[0]);
+      int32_positions.resize(static_cast<size_t>(length));
+      std::transform(
+          pos.data<int64_t>(),
+          pos.data<int64_t>() + length,
+          int32_positions.begin(),
+          [](int64_t p) { return static_cast<int32_t>(p); });
+      positions = int32_positions.data();
       break;
     default:
       throw std::runtime_error(
@@ -332,7 +348,8 @@ inline void exec_update_and_attend(
   }
   AttendSpec spec = st.cache->update_and_fetch(
       *n.layer_id,
-      position,
+      positions,
+      length,
       st.const_tensor_ref(n.k),
       st.const_tensor_ref(n.v),
       s);

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

backends/mlx/runtime/MLXSequenceCache.h

@@ -8,14 +8,14 @@
 
 #pragma once
 
-#include <algorithm>
 #include <optional>
 #include <stdexcept>
 #include <string>
 #include <vector>
 
 #include "MLXCache.h" // AttendSpec, MLXCache
 #include "MLXExecutor.h" // resolve_dtype
+#include "MLXPool.h" // Pool
 
 #include <executorch/extension/llm/cache/cache.h>
 #include <executorch/extension/llm/cache/sequence_cache.h>
@@ -26,97 +26,6 @@ namespace mlx {
 
 namespace cache = ::executorch::extension::llm::cache;
 
-// Per-layer K or V store, SDPA-major [1, H, slots, D] (cells on axis 2). The
-// planner hands down physical runs (it has already applied any ring modulo), so
-// the pool is layout-agnostic: flat and ring differ only in how many slots the
-// layer asks for and how many runs 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 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 > max_slots_) {
-      throw std::runtime_error("Pool::write: run out of bounds");
-    }
-    if (static_cast<int>(update.shape(2)) != run.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(run.start + run.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, run.start, 0},
-        ::mlx::core::Shape{1, H, run.start + run.len, D},
-        s);
-  }
-
-  // The run's cells, [start, start+len). Ring reads start mid-pool, so the run
-  // start matters here as much as it does for a write.
-  Tensor read(const cache::Run& run, StreamOrDevice s) const {
-    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()) {
-      throw std::runtime_error("Pool::read: run out of bounds");
-    }
-    return ::mlx::core::slice(
-        buf_,
-        ::mlx::core::Shape{0, 0, run.start, 0},
-        ::mlx::core::Shape{1, H, run.start + run.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_;
-};
-
 // 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
@@ -172,7 +81,8 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
 
   AttendSpec update_and_fetch(
       int layer,
-      int position,
+      const int32_t* positions,
+      int length,
       const Tensor& k,
       const Tensor& v,
       StreamOrDevice s) override {
@@ -181,7 +91,8 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
     }
     const int T = static_cast<int>(k.shape(2)); // BHSD: seq axis is 2
 
-    std::optional<cache::SeqStepPlan> p = this->plan(layer, position, T);
+    std::optional<cache::SeqStepPlan> p =
+        this->plan(layer, run_start(positions, length, T), T);
     if (!p) {
       throw std::runtime_error(
           "update_and_fetch: step exceeds capacity or invalid layer");
@@ -216,6 +127,24 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
   }
 
  private:
+  // A sequence cache holds one run of one sequence, so the step is described by
+  // where it starts; the remaining positions carry no information beyond
+  // confirming that. A step this cache cannot represent is refused rather than
+  // silently stored at the wrong positions.
+  static int run_start(const int32_t* positions, int length, int T) {
+    if (length != T) {
+      throw std::runtime_error(
+          "update_and_fetch: one position per query token expected");
+    }
+    for (int i = 1; i < length; ++i) {
+      if (positions[i] != positions[0] + i) {
+        throw std::runtime_error(
+            "update_and_fetch: sequence cache needs a contiguous run");
+      }
+    }
+    return positions[0];
+  }
+
   // 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.
@@ -226,15 +155,16 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
       const Tensor& update,
       StreamOrDevice s) {
     if (n == 1) {
-      pool.write(runs[0], update, s);
+      pool.write(runs[0].start, runs[0].len, update, s);
       return;
     }
     const int H = static_cast<int>(update.shape(1));
     const int D = static_cast<int>(update.shape(3));
     int off = 0;
     for (int i = 0; i < n; ++i) {
       pool.write(
-          runs[i],
+          runs[i].start,
+          runs[i].len,
           ::mlx::core::slice(
               update,
               ::mlx::core::Shape{0, 0, off, 0},
@@ -250,12 +180,12 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
   static Tensor
   read_runs(const Pool& pool, const cache::Run* runs, int n, StreamOrDevice s) {
     if (n == 1) {
-      return pool.read(runs[0], s);
+      return pool.read(runs[0].start, runs[0].len, s);
     }
     std::vector<Tensor> parts;
     parts.reserve(static_cast<size_t>(n));
     for (int i = 0; i < n; ++i) {
-      parts.push_back(pool.read(runs[i], s));
+      parts.push_back(pool.read(runs[i].start, runs[i].len, s));
     }
     return ::mlx::core::concatenate(parts, 2, s);
   }