[MLX] Add off-graph KV-cache flat runtime

[kiymetakdemir]

Summary

Adds the MLX runtime for the off-graph KV cache: a flat (full-history) cache plus the update_and_attend op handler, wiring the neutral SequenceCache/FlatPolicy bookkeeping to real MLX tensors. The graph carries a cache-agnostic UpdateAndAttendNode — no cache operand, since the cache is off-graph runtime state keyed by layer_id. At runtime the handler asks the cache for the K/V window and mask kind (AttendSpec) and calls SDPA: the cache owns KV byte movement + attention semantics, the handler owns q/scale + the SDPA call. Single-sequence, fp16, full-history; the per-query mask stays in the backend as MLX's fused "causal"/none (no materialized tensor).

Files

• backends/mlx/serialization/schema.fbs — UpdateAndAttendNode (q,k,v,position,out; layer_id, scale, causal) appended to the OpNode union.
• backends/mlx/runtime/MLXCacheImpl.h — op-facing MLXCacheImpl interface + AttendSpec (the K/V window + mask kind None/Causal/Explicit; window (SWA) and score/softcap are deferred.
• backends/mlx/runtime/MLXFlatCache.h — FlatPool + MLXFlatCache.
• backends/mlx/runtime/MLXExecutor.h — ExecutionState gains a per-session MLXCacheImpl* cache (survives reset(), mirrors mutable_buffers).
• backends/mlx/runtime/MLXInterpreter.h — exec_update_and_attend handler (fetch AttendSpec → SDPA) + the OpCode::UPDATE_AND_ATTEND dispatch case.
• backends/mlx/test/mlx_flat_cache_test.cpp, backends/mlx/test/CMakeLists.txt — op-level parity test.

Testing

C++ op-level test: drives MLXFlatCache::update_and_fetch directly and attends the returned AttendSpec exactly as the handler does, comparing to MLX SDPA over the full K/V history the cache should have assembled, prefill (Causal), decode (None), and the capacity-reject path. Runs on Apple Silicon (MLX SDPA needs the Metal backend). Built and ran via the standard flow, passes:

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 25 Pending

As of commit 83085d9 with merge base e0c610a ():
💚 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]

Didn't the op have an output dtype?

[metascroy]

kv here is just update?

[kiymetakdemir]

Yes it is.

[metascroy]

Why is capacity a property of write?

[kiymetakdemir]

Fixing this.

[metascroy]

Is return type Tensor or ::mlx::array?

[kiymetakdemir]

Tensor is the runtime's existing alias for ::mlx::core::array (in MLXExecutor.h), so I used it as in there.

[kiymetakdemir]

Recent changes:

1. Op contract; dropped causal, added out_dtype
2. Allocate at construction
3. Bounds assertion on the write position
4. Storage dtype from config + cast-if-differ
   Note: I kept capacity, it is the max cap, behaves like static now, later will allocate initial capacity and have dynamic grow.

[metascroy]

Raise if dtype not specified.

[metascroy]

I'm quite worried about this. It breaks async evalution.

Can we lift this logic to exec_update_and_attend at least and pass an int down. I think it'll be easier to correct at that level if needed.

[metascroy]

Use GTEST

[metascroy]

We also need something to run the c++ test in CI. You might need to modify mlx.yml

[kiymetakdemir]

Okay, added!

[metascroy]

construct shape directly? Why go through vector?

[kiymetakdemir]

I guess it is a leftover from lazy first write, fixing it.

[metascroy]

Curious: what is the rational for the indirection here? Why have MLXCacheImpl?

Isn't MLXFlatCache an impl itself?

[metascroy]

If it's required field, why make it optional?

[metascroy]

Something about this seems a bit off. In sequence cache we have have ring/flat/etc as different policies on the same class b/c different layers of the cache have different policies.

But in the implementation here, you've made flatcache as the top-level (instead of something like mlxsequencecache), and ignored the layer-by-layer policy on cfg, and assumed all are flat.

How do you intend to extend this to ring in the next PR?

[kiymetakdemir]

Right, now I restructured to address ring as well. Renamed MLXFlatCache to MLXSequenceCache and the constructor reads lc.policy per layer. Does it look like the right structure?

[metascroy]

What is the point of running SDPA twice here? Don't you just need to test the spec0.K = k0 etc?

[kiymetakdemir]

True, changing it.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 5m 40s —— View job

---

Code Review

Reviewed the full diff against origin/main. This is a clean, well-scoped runtime addition and the neutral SequenceCache/FlatPolicy bookkeeping is nicely separated from the MLX byte layer. The prior review threads (out_dtype, allocate-at-construction, bounds assert, required kv_dtype, GTest, CI wiring) are all addressed. A few things worth a look, roughly by importance.

Correctness / behavior

1. scale has a silent-wrong default (schema.fbs:161). scale: float; defaults to 0.0. If the exporter ever omits it, fast::scaled_dot_product_attention runs with scale 0 → uniform attention over V, with no error — a very hard-to-diagnose numerical bug. q/k/v/position/out are all (required); scale is effectively required too. Consider making it (required) or documenting/asserting a sane default in the handler. Same reasoning the team already applied to kv_dtype.

2. Per-layer host sync remains (MLXInterpreter.h:311–313). You lifted the eval(pos)/read into exec_update_and_attend and pass an int down, which resolves @metascroy's async-eval concern at the interface level. But the eval(pos) still fires once per UpdateAndAttendNode, i.e. once per layer per step — N_layers device→host stalls per decode token, even though position is identical across all layers of a step. The comment already flags the host-constant path as the real fix; worth noting the cost is per-layer, not per-step, so it's N× worse than it may read. Fine to defer, but a comment quantifying that would help.

3. FlatPool graph growth across steps (MLXFlatCache.h:57). buf_ = slice_update(buf_, ...) chains a new lazy node on top of the prior buf_ each write. This stays bounded only because the SDPA output is eval'd each step (materializing the read slice, and transitively buf_). That's an implicit invariant — if a future caller ever runs update_and_fetch without eval'ing the result, the graph grows unboundedly per step. A one-line note on the eval dependency would make the contract explicit.

Design (raised by @metascroy — I concur)

4. MLXFlatCache as top-level vs. per-layer policy (MLXFlatCache.h:79–98). The neutral core is explicitly per-layer-policy (layer_to_policy_, flat/ring dedup), but MLXFlatCache allocates a FlatPool for every layer and hard-rejects any 2-run plan in update_and_fetch (lines 120–123). So a ring layer doesn't silently corrupt — it throws — which is the safe failure. But the class name and structure bake in "all flat." For the ring PR you'll likely want an MLXSequenceCache holding a unique_ptr<Pool> per layer chosen from lc.policy.kind, rather than growing MLXFlatCache. Might be worth landing the neutral shape now to avoid a rename churn later.

5. CacheBase* → MLXCacheImpl* cross-cast (MLXCacheImpl.h:31–33). The header comment says "cross-cast from the registry's CacheBase*… no dynamic_cast/RTTI." SequenceCache(→CacheBase) and MLXCacheImpl are sibling bases of MLXFlatCache; recovering one base subobject from the other genuinely requires either dynamic_cast or a static_cast through the concrete type. Since the install/registry code isn't in this PR, I can't see how the no-RTTI cross-cast is done — worth confirming that path exists and compiles before this is reachable (st.cache is currently never assigned anywhere in the tree, so the op is dead until a follow-up wires it).

Minor

6. to_shape(std::vector<int>{...}) (MLXFlatCache.h:37, and the test). As @metascroy noted, you can construct ::mlx::core::Shape{1, H, capacity, D} directly and skip the intermediate std::vector + to_shape conversion.

7. AttendSpec::Mask → mask_mode mapping (MLXInterpreter.h:324). kind == Causal ? "causal" : "" collapses None and Explicit to the same "" mode, relying on spec.mask being nullopt vs. set to distinguish them. It's correct today, but an explicit switch on kind would be more robust once Explicit is actually produced (and would catch a future Explicit-with-nullopt mistake).

8. Prefill test runs SDPA twice (mlx_flat_cache_test.cpp:94–105). For the prefill case spec0.K/spec0.V are just the first write, so EXPECT_TRUE(allclose(spec0.K, k0, ...)) would test the assembled-window contract more directly and cheaply. The double-SDPA comparison earns its keep in DecodeAttendsFullHistory (spec vs. concatenated history) but is redundant here. Not blocking.

Nothing here is a blocker for a scoped runtime PR. #1 (scale default) and #5 (the cross-cast reachability) are the two I'd want confirmed before the op goes live.
• branch kv-cache-mlx-flat

[metascroy]

I don't think this comment is accurate

[metascroy]

@claude review this code

[claude]

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

---

Code Review

Re-reviewed the full diff (HEAD~3..HEAD, base 43cb2b2). Since my last pass the restructure landed well: MLXFlatCache → MLXSequenceCache reads lc.policy per layer, ring is refused at construction rather than mis-served, kv_dtype is a plain (required) int, scale is float = null with a loud handler check, the mask switch is exhaustive, Pool::write's graph-growth invariant is documented, and the per-node host-sync cost is called out at MLXInterpreter.h:310-314. The op-level test now compares the assembled window directly instead of double-SDPA. Good iteration.

A few remaining observations, none blocking:

1. No bounds check on cfg.layers.size() vs n_layers (MLXSequenceCache.h:113-114). The ctor broadcasts when layers.size() == 1, else indexes cfg.layers[l] for l ∈ [0, n_layers). If a caller ever passes a layers vector whose size is neither 1 nor n_layers (say a partial per-layer list), cfg.layers[l] reads out of bounds with no diagnostic. A single guard — layers.size() == 1 || layers.size() == n_layers — would turn a UB read into a clear error, and matches the contract the comment in cache.h:123-124 already states. Fix this →

2. layer_id: int32; in the schema (schema.fbs:161) has no (required) and no null default, so it silently defaults to 0. Unlike scale, layer 0 is a legitimate value, so an omitted layer_id won't fail loudly — it'll route every layer's KV to pool 0 and corrupt the cache quietly. Worth the same treatment scale/kv_dtype got (mark (required) or = null + assert), since it's the same class of exporter-omission bug the team already hardened against.

3. commit() is called once per layer inside update_and_fetch (MLXSequenceCache.h:153). SequencePlanner's contract (cache.h:83-85) says commit is idempotent and "calling it once per step — not per layer — suffices," so this is correct, just redundant N_layers times per step. Harmless today; only flagging because the interface comment implies the intended call site is per-step (the runner), and once ring lands with a real per-step position you'll want to be sure the double-accounting stays a genuine no-op.

4. Minor — to_shape(std::vector<int>{...}) still in the test (mlx_sequence_cache_test.cpp:69). The runtime Pool ctor now builds ::mlx::core::Shape{...} directly (good); the test's randn helper still routes through the vector. Cosmetic, and test-only.

Nothing here blocks the merge. #1 and #2 are the two cheap hardening steps I'd fold in — both close silent-corruption paths in the same spirit as the scale/kv_dtype fixes already made.
• branch kv-cache-mlx-flat