[MLX] Add off-graph KV cache ring runtime

[kiymetakdemir]

Summary

Adds ring (sliding-window) layers to the MLX KV cache, so a model can mix them with flat layers per layer. The window is expressed declaratively: AttendSpec gains a window on Causal (the mask_mod axis the header reserved), so the cache states the semantic and the handler materializes whatever the backend can't do natively. MLX SDPA has only ""/"causal"/"array" — no sliding-window mode — so the handler builds a banded bool mask, but only when the window is narrower than the span. A backend that supports windows natively would just forward the integer.

Files

• backends/mlx/runtime/MLXCache.h — AttendSpec gains std::optional window (Causal only; unset = unbounded history).
• backends/mlx/runtime/MLXInterpreter.h — window_causal_mask(T, S, window) builds the [1,1,T,S] band from two arange compares (on-device, no host sync); the
  handler uses it only when window < S, else plain "causal".
• backends/mlx/runtime/MLXSequenceCache.h — ring layers size their pool at window + max_write - 1 and are allocated outright rather than grown (a ring is already bounded, and growing it would complicate the modulo the planner applied); per-layer window is recorded and declared on multi-token steps; write_runs/read_runs scatter and gather the 1–2 runs a wrapping step produces; the construction-time ring rejection is removed.
• backends/mlx/test/mlx_sequence_cache_test.cpp — four cases.

Testing

New: the mask is a band (checked against a hand-written 3×5 matrix, plus a window covering the span degenerating to causal); a ring decode past its window evicts the oldest and needs no mask; a multi-token step declares its window while a flat layer declares none; and a step whose runs wrap the ring is scattered and rejoined in logical order; the physical layout is out of order by then, so this pins the wrap handling.

cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON
cmake --build cmake-out --target mlx_sequence_cache_test
ctest --test-dir cmake-out -R mlx_sequence_cache --output-on-failure

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 94187b1 with merge base e8feb9e ():
💚 Looks good so far! There are no failures yet. 💚

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]

Isn't there a field for mask on attendsepc?

Should the mask be constructed in the cache and handed back?

[kiymetakdemir]

Mlx doesnt have a window mode but for backends that takes window like flash attention it could just pass window parameter instead of materializing the mask

[metascroy]

I'm a bit confused. On AttendSpec, you have std::optional mask; // Explicit only

And you just overwrite that value here?

Why can the logic in window_causal_mask not live in MLXSequanceCache.h and attach a mask there?

[kiymetakdemir]

Actually, you're right. I was treating AttendSpec as it is a cross-backend contract thinking that it will declare the window and each backend can render it but each backend will have its own cache class anyway. I am fixing this.

[metascroy]

Creating this mask is not free, are you creating it once per layer IIUC

Per decode step, we only need to plan once per policy per execute step. Is that how your code is set up?

Much of this will become more apparent when we test in real model.

[kiymetakdemir]

Decode currently doesn't need a mask. But no, we don't plan once per policy. Policies aren't duplicated, but update_and_fetch is per layer, so we call plan() for each layer per step. I can memoize the plan and mask per policy index, invalidated when (position, T) changes.

[metascroy]

Why doesn't decode need a mask? Isn't the mask how you tell the ring what to attend to?

[kiymetakdemir]

The selection happens in the read, the planner returns the retained window as physical runs, and we read exactly those. For example, W=4, max_write=1 (ring of 4 slots), decoding at position 10. RingPolicy::plan(10, 1) gives rstart = 7, span 4, which wraps into two runs, {start:3, len:1} and {start:0, len:3}. We read those and concatenate, resulting positions 7, 8, 9, 10 in logical order.

We need a mask during prefill: a step with T > 1 reads the union of its queries' visibilities.

[metascroy]

Even in prefill, I'm not sure we'll want to generate a mask per layer.

Stamping to unblock, but consider doing plan once per execute per policy, not per layer.

I guess the next PR will be E2E enablement, and we can start comparing perf.

[kiymetakdemir]

Thanks! Do you think memoization would make sense for this problem? Keeping the plan (and mask) per policy index and invalidating when (position, T)?

[metascroy]

Even in prefill, I'm not sure we'll want to generate a mask per layer.

Stamping to unblock, but consider doing plan once per execute per policy, not per layer.

I guess the next PR will be E2E enablement, and we can start comparing perf.