[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 c921c59 with merge base 849f9d8 ():
💚 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