[MLX] Add off-graph KV-cache flat runtime by kiymetakdemir · Pull Request #21501 · pytorch/executorch

kiymetakdemir · 2026-07-30T19:13:59Z

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 · 2026-07-30T19:14:04Z

🔗 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, 17 Pending

As of commit 011bd32 with merge base 43cb2b2 ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2026-07-30T19:14:48Z

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 · 2026-07-30T19:27:34Z

    causal: bool = false;
 }

+table UpdateAndAttendNode {


Didn't the op have an output dtype?

metascroy · 2026-07-30T19:31:02Z

+// write is a slice_update run, a read a [0, len) slice. H/D inferred on write.
+class FlatPool {
+ public:
+  void write(int start, const Tensor& kv, int capacity, StreamOrDevice s) {


kv here is just update?

metascroy · 2026-07-30T19:31:52Z

+// write is a slice_update run, a read a [0, len) slice. H/D inferred on write.
+class FlatPool {
+ public:
+  void write(int start, const Tensor& kv, int capacity, StreamOrDevice s) {


Why is capacity a property of write?

Fixing this.

metascroy · 2026-07-30T20:08:48Z

+    buf_ = ::mlx::core::slice_update(*buf_, kv, to_shape(lo), to_shape(hi), s);
+  }
+
+  Tensor read(int len, StreamOrDevice s) const {


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

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

kiymetakdemir · 2026-07-30T23:36:11Z

Recent changes:

Op contract; dropped causal, added out_dtype
Allocate at construction
Bounds assertion on the write position
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 · 2026-07-30T23:40:30Z

+      : cache::SequenceCache(cfg) {
+    const ::mlx::core::Dtype dt = cfg.kv_dtype
+        ? resolve_dtype(static_cast<int8_t>(*cfg.kv_dtype))
+        : ::mlx::core::float16;


Raise if dtype not specified.

metascroy · 2026-07-30T23:47:05Z

+  // position vector is tiny). Cast to int32 so an int64 position also works.
+  static int read_start(const Tensor& position, StreamOrDevice s) {
+    Tensor p = ::mlx::core::astype(position, ::mlx::core::int32, s);
+    ::mlx::core::eval(p);


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 · 2026-07-30T23:48:46Z

+#include "MLXFlatCache.h"
+
+#include <mlx/mlx.h>
+


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

Okay, added!

kiymetakdemir requested review from SS-JIA, digantdesai and metascroy July 30, 2026 19:13

kiymetakdemir requested review from kirklandsign and larryliu0820 as code owners July 30, 2026 19:14

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 30, 2026

metascroy reviewed Jul 30, 2026

View reviewed changes

kiymetakdemir force-pushed the kv-cache-mlx-flat branch 2 times, most recently from 88761ae to 9321153 Compare July 30, 2026 23:23

kiymetakdemir requested a review from mergennachin as a code owner July 30, 2026 23:23

metascroy reviewed Jul 30, 2026

View reviewed changes

kiymetakdemir force-pushed the kv-cache-mlx-flat branch from 9321153 to 4d88209 Compare July 30, 2026 23:47

metascroy reviewed Jul 30, 2026

View reviewed changes

kiymetakdemir force-pushed the kv-cache-mlx-flat branch 2 times, most recently from ad06aae to b27b89c Compare July 31, 2026 00:06

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

011bd32

kiymetakdemir force-pushed the kv-cache-mlx-flat branch from b27b89c to 011bd32 Compare July 31, 2026 00:08

		#include "MLXFlatCache.h"

		#include <mlx/mlx.h>

Uh oh!

Conversation

kiymetakdemir commented Jul 30, 2026

Uh oh!

pytorch-bot Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

⏳ No Failures, 17 Pending

Uh oh!

github-actions Bot commented Jul 30, 2026

This PR needs a release notes: label

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiymetakdemir commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Jul 30, 2026 •

edited

Loading

This PR needs a `release notes:` label

kiymetakdemir commented Jul 30, 2026 •

edited

Loading