[MLX] Per-token positions and a shared pool

[kiymetakdemir]

Summary
Prepares the MLX byte layer to take a second cache implementation. No behaviour changes.

MLXCache::update_and_fetch takes per-token positions instead of a run start.
The position tensor already carries one entry per query token and the interpreter was reading element [0] and discarding the rest, so this stops truncation. Pool moves to its own header and takes (start, len) instead of cache::Run, so it no longer depends on the neutral cache header.

Files

• backends/mlx/runtime/MLXCache.h — take per-token positions.
• backends/mlx/runtime/MLXInterpreter.h — eval and the dtype switch keep their shape; int32 now passes the tensor's own buffer with no copy, only int64 is narrowed, and pos.size() is checked against k.shape(2).
• backends/mlx/runtime/MLXPool.h — new; Pool, unchanged apart from the signature.
• backends/mlx/runtime/MLXSequenceCache.h — run_start, and the five Pool call sites.
• backends/mlx/test/mlx_sequence_cache_test.cpp — a step() helper builds the run from k.shape(2)

Testing

ctest --test-dir cmake-out -R mlx_sequence_cache

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22037

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (6 Unrelated Failures)

As of commit 38e6b21 with merge base bd8b2be ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

This is just a move of the code that used to be in backends/mlx/runtime/MLXSequenceCache.h, right?

[metascroy]

This memory isn't owned by us. What prevents pointer being bad?

[metascroy]

Why the raw pointer? Can we pass a const vec ref?

[metascroy]

Maybe always cp to positions vec to control the memory?

[metascroy]

Compare against n.q?

[metascroy]

This is the code block where you had a perf regressin before. Did you verify that these changes don't cause regression?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 52s —— View job

---

Code review

Reviewed the diff against origin/main. This is a clean, well-commented refactor and the "no behaviour change" claim mostly holds. Answers to the inline questions plus a couple of findings below.

Answers to your inline questions

MLXPool.h:1 — "just a move?" Yes. Pool is byte-for-byte identical to what was in MLXSequenceCache.h except the write/read signatures change from const cache::Run& to (int start, int len), and the doc comment swaps "planner/flat and ring" for "caller/policies". That's the intended point of the PR — dropping the dependency on the neutral cache header.

MLXInterpreter.h:333 — "this memory isn't owned by us, what prevents a bad pointer?" It's safe as written, but only by a subtle invariant worth a comment:

• auto pos = st.const_tensor_ref(n.position) copies the MLX array handle (refcount bump on the shared buffer), so pos keeps the buffer alive for the whole function.
• positions = pos.data<int32_t>() points into that buffer, and it's consumed synchronously inside update_and_fetch before pos goes out of scope.

So lifetime is fine. The real latent hazard is contiguity, not ownership: data<int32_t>() + linear iteration assumes a flat, row-contiguous buffer. eval(pos) materializes the array but doesn't guarantee it isn't a strided/offset view. In practice the position tensor is a fresh 1-D int array so this holds, but a one-line guard would make it robust:

assert(pos.flags().row_contiguous);

MLXInterpreter.h:351 — "why the raw pointer? can we pass a const std::vector&?" The raw pointer is what avoids a copy on the int32 hot path — it hands update_and_fetch the array's own buffer directly. A const std::vector<int32_t>& would force materializing a vector for every int32 step, which is the regression you're worried about at :336. If you want a safer signature without the copy, a (const int32_t*, int) pair is exactly a hand-rolled span; std::span<const int32_t> (C++20) would carry the length together with the pointer and read better, if the codebase allows it.

MLXInterpreter.h:329 — "maybe always cp to a positions vec to control the memory?" That trades the ownership subtlety above for a copy on the common int32 path. Given the perf note in the file, I'd keep the zero-copy int32 path and instead add the contiguity assert + a short comment stating "pos (the array handle) owns the buffer that positions points into; do not let it go out of scope before update_and_fetch returns." Cheaper and documents the invariant.

MLXInterpreter.h:324 — "compare against n.q?" Agreed, this is the more correct check. positions is one entry per query token, so the natural invariant is length == q.shape(2), not k.shape(2). They're equal for standard self-attention (new K and Q share the seq axis), so behaviour is identical today — but the assertion should express the thing it actually means. If you keep k, at least update the error string, which currently says "one entry per query token" while checking K.

MLXInterpreter.h:336 — "you had a perf regression here before; verified?" This branch only runs for int64 position tensors (the resize + std::transform narrowing). The int32 path is now genuinely zero-copy — strictly better than before. So the previous regression path (a per-layer copy on the hot dtype) is avoided. Worth confirming the exported models actually feed int32 here so the int64 copy isn't silently on the hot path; if they can be int64, that per-step allocation is the thing to watch.

Other observations

1. Double validation of length == T. MLXInterpreter.h:324 checks length == k.shape(2), then MLXSequenceCache::run_start (MLXSequenceCache.h:135) re-checks length != T and throws. Harmless redundancy, but the interpreter check makes the cache-side one unreachable via the normal path. That's fine as defense-in-depth since the cache is a public interface, just noting it's intentional-looking.

2. Test coverage is good. NonContiguousOrMiscountedPositionsThrow covers the gap, two-sequence, and short-run cases, and the step() helper cleanly threads std::iota-generated positions through all existing tests. The PoolHonorsRunStart regression guard is preserved.

3. Nit: MLXInterpreter.h:342 — positions = int32_positions.data() is correct only because int32_positions outlives the update_and_fetch call; same-scope, so fine.

Nothing blocking. The one change I'd actually make is the k.shape(2) → q.shape(2) check (or fix the error string), plus a one-line comment/assert nailing down the buffer-lifetime + contiguity invariant at :329–333.
· branch mlx-byte-layer-prep