[MLX] Add the cell KV-cache layout to the MLX runtime

[kiymetakdemir]

Summary

MLXCellCache turns a CellStep into a scatter over the cell axis and materialises the step's mask bits as an explicit SDPA mask, one row per query token. Pool grows a matching write_cells next to its existing contiguous
write, keeping the pool layout-agnostic: layouts differ only in which slots they ask for. The layout is registered as the "cell" builder alongside "seq", which is how a runner selects one.

Files

• backends/mlx/runtime/MLXPool.h — write_cells, the scatter primitive
• backends/mlx/runtime/MLXCellCache.h — new; the CellStepper face over Pool
• backends/mlx/runtime/MLXBackend.cpp — registers the "cell" builder
• backends/mlx/test/mlx_cell_cache_test.cpp — new; 11 tests
• backends/mlx/test/CMakeLists.txt, .github/workflows/mlx.yml — build and run it

Test

mlx_cell_cache_test, 11 tests. Covers cell claiming on prefill and decode, the mask isolating sequences, a freed cell
refilling below a live one (the case no contiguous write can express, so the only real exercise of the scatter), a windowed layer hiding older cells, pool growth preserving written cells, storage-dtype casts, the step-verb contract, capacity refusal, and the registry lookup by (backend_id, kind).

Added to the existing test-mlx CI job.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Pending

As of commit cec6bf6 with merge base f3982af ():

NEW FAILURE - The following job has failed:

• pull / android / build-android (gh)

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]

@claude review this code

[claude]

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

---

Code Review — [MLX] Cell KV-cache layout

• Read the diff and new files (MLXCellCache.h, MLXPool::write_cells, tests, registration)
• Cross-check against the neutral CellCache contract and the sibling MLXSequenceCache
• Analyze correctness, dtype handling, bounds, and test coverage

---

Overall this is clean, well-scoped, and reads like its neighbors — the write_cells scatter mirrors write, the mask handling matches MLXSequenceCache, and the comments carry the non-obvious context in the house style. Test coverage is genuinely good; FreedCellsRefillBelowLiveOnes is the right stress case for the scatter. A few things worth a look.

1. write_cells doesn't explicitly cast to the storage dtype (the one substantive concern)

MLXPool::write casts the update to dtype_ before the slice_update so the buffer's dtype stays stable:

const Tensor u = update.dtype() == dtype_ ? update : ::mlx::core::astype(update, dtype_, s);

write_cells (MLXPool.h:91) skips this and leans on put_along_axis instead — the comment even claims it "casts the update to the storage dtype", but there's no astype there:

buf_ = ::mlx::core::put_along_axis(buf_, /*indices*/, update, 2, s);

If put_along_axis type-promotes (buf_, update) rather than casting update down to buf_'s dtype, then feeding f32 K/V into a bf16/f16 pool silently turns buf_ into f32. That corrupts the storage invariant, and the next maybe_grow concatenates a dtype_ (bf16) pad onto an f32 buffer. StorageDtypeDiffersCastsOnWrite asserts spec.K.dtype() == bfloat16, so CI on Apple Silicon will tell us which way MLX resolves this — but relying on put_along_axis's promotion semantics for a correctness invariant is fragile. I'd cast explicitly for parity with write, so the invariant holds regardless of MLX's promotion rules and the comment becomes true:

const Tensor u = update.dtype() == dtype_ ? update : ::mlx::core::astype(update, dtype_, s);
buf_ = ::mlx::core::put_along_axis(buf_, /*indices*/, u, 2, s);

Fix this →

2. Duplicate cell indices are undefined for the scatter — worth an assertion or comment

put_along_axis has unspecified behavior when two indices collide. Correctness here rests entirely on the neutral layer guaranteeing distinct cells per step (lowest_free per token). That holds today, but it's an invisible cross-layer invariant — the scatter would silently pick a winner if it were ever violated. A one-line comment noting "cells are distinct per step (guaranteed by place())" next to the put_along_axis call would document why no dedup is needed. Not blocking.

3. read(0, read_len) safety depends on monotonic pool growth — correct, but subtle

update_and_fetch reads [0, read_len) where read_len == used_end, but write_cells only grows the pool to high = max(cell)+1 for this step's cells. When a step refills a low cell while higher cells stay live (exactly FreedCellsRefillBelowLiveOnes), read_len > high. This is safe only because the pool never shrinks and every occupied cell in [0, used_end) was grown-to by some earlier write_cells. That reasoning is correct but non-obvious; a brief note in MLXPool::read or at the call site would help a future reader (the existing "grows toward max_slots on demand" comment doesn't quite cover the read-past-this-step's-high case).

Minor / nits

• MLXPool.h:135 — the maybe_grow comment still says "write() already bounds needed by max_slots_". write_cells now also bounds it (via the per-cell cell >= max_slots_ check), so the comment is slightly stale; consider "callers bound needed by max_slots_".
• IllFormedStepsThrow and StepPastCapacityIsRefused are nicely targeted. One gap: no test exercises Pool::write_cells's own guards (e.g. update.shape(2) != cells.size()), since update_and_fetch validates lengths before reaching the pool. Low value given the upstream check, but it's the only branch in the new code without direct coverage.

Nothing here blocks — item 1 is the only correctness question, and CI will resolve the dtype behavior; making the cast explicit removes the dependency on MLX's promotion semantics either way.
· branch mlx-cell-cache

[metascroy]

Did I see a similar helper in one of your previous PRs? If so, let's factor out into test/utils.h and reuse

[metascroy]

Is capacity here capacity across all sequences? So even for windowed layers, total capacity can exceed the window size?

[kiymetakdemir]

Capacity is for total tokens across all sequences. Yes we have the same pool size for windowed layers because we have a global cell table not a table per policy. place() claims one cell index per token per forward, and every layer's K/V pools are addressed by the same index. Then we mask according to window for each sequence.

[kiymetakdemir]

On claude's review; put_along_axis casts the update to the destination's dtype inside scatter_axis (mlx/ops.cpp:3642). We have a test case StorageDtypeDiffersCastsOnWrite writes fp32 into a bf16 pool and passes.

[metascroy]

On claude's review; put_along_axis casts the update to the destination's dtype inside scatter_axis (mlx/ops.cpp:3642). We have a test case StorageDtypeDiffersCastsOnWrite writes fp32 into a bf16 pool and passes.

I think Claude's point is to just be explicit about it, rather than rely on type promotion (which I agree with).

Stamping b/c I think it looks good, but consider the explicit cast

[metascroy]

Why are ring_config/flat_config test utils?

Couldn't the contructor on CacheConfig do what they do?

[kiymetakdemir]

CacheConfig doesn't have a constructor, it's an aggregate