Add a batched-sequence KV cache layout and reorganize the control faces by kiymetakdemir · Pull Request #22662 · pytorch/executorch

kiymetakdemir · 2026-09-10T00:40:38Z

Summary

A third KV cache layout: many sequences, each over its own private history, sharing one capacity. A step still arrives flat on the token axis, but is answered per sequence. Against the cell layout on llama-1B/M4 Max it crosses over around B=8 and is 17–36% faster at B=32; below that it costs 2–5%.

Files

extension/llm/cache/cache.h — CacheControl down to capacity/clear; can_extend(n) on SequenceControl; seq_rm(id), rewind(id, new_len), can_admit(id, n), max_seqs() on BatchControl.

extension/llm/cache/cell_cache.{h,cpp} — the same interface on the incumbent: seq_rm/rewind over a shared drop_from, can_admit keyed off next_pos, max_seqs = kMaxSeqs.

extension/llm/cache/sequence_cache.h — can_extend gains the max_context term; copy constructor and length() so a batched row can fork.

extension/llm/cache/cache_registry.h, CMakeLists.txt, README.md — kind::kBatchedSequence, the new .cpp in extension_llm_cache, the renamed verb in the example.

backends/mlx/runtime/MLXCache.h — attend(layer, positions, q, k, v, scale, s) becomes the single virtual; a free attend(AttendSpec, q, scale, s) holds the dtype match and mask-mode switch, so a layout answering per sequence needn't materialise every sequence's history first.

backends/mlx/runtime/MLXSequenceCache.h — gains attend(), which runs update_and_fetch and hands the spec to the shared SDPA helper.

backends/mlx/runtime/MLXBatchedSequenceCache.h — the new layout. Implements BatchControl over std::map<seq_id, MLXSequenceCache>, so each sequence's bookkeeping and pools are the single-sequence cache that already existed.

backends/mlx/runtime/MLXCellCache.h, MLXInterpreter.h, MLXBackend.cpp — cell gets attend over its update_and_fetch, exec_update_and_attend shrinks to one call, the third builder is registered.

backends/mlx/examples/llm/run_llm_batched.cpp — --cache_kind, defaulting to batched-cell.

extension/llm/cache/test/cache_test.cpp, backends/mlx/test/mlx_batched_sequence_cache_test.cpp, backends/mlx/test/CMakeLists.txt, mlx_cell_cache_test.cpp — neutral tests for spans, placement, forks, context and admission; op-level parity test whose oracle is an MLXSequenceCache fed the same tokens.

extension/llm/cache/reference_cache.py, test_update_and_attend.py — the eager oracle tracks the same interface: max_context, can_admit(seq_id, n), and a next_pos on the cell reference, which didn't have one.run-length-encoding ids into spans, place_step planning each span against that sequence's own SequenceCache, seq_clone copying a prefix.

extension/llm/cache/cell_cache.{h,cpp} — the same interface on the incumbent: seq_rm/rewind over a shared drop_from, can_admit keyed off next_pos rather than seq_len (positions need not be consecutive there), max_seqs = kMaxSeqs.

extension/llm/cache/sequence_cache.h — can_extend gains the max_context term; copy constructor and length() so a batched row can fork.

extension/llm/cache/cache_registry.h, CMakeLists.txt, README.md — kind::kBatchedSequence, the new .cpp in extension_llm_cache, the renamed verb in the example.

backends/mlx/runtime/MLXCache.h — attend(layer, positions, q, k, v, scale, s) becomes the single virtual; a free attend(AttendSpec, q, scale, s) holds the dtype match and mask-mode switch, so a layout answering per sequence needn't materialise every sequence's history first.

backends/mlx/runtime/MLXSequenceCache.h — mostly relocation: write_runs, read_runs, LayerPools, make_pools, spec_for from private members to free functions, bodies byte-identical; three pool members collapse into one LayerPools; attend() added over the existing update_and_fetch. Read with --color-moved=zebra. Flat pools now cap at min(capacity, max_context).

backends/mlx/runtime/MLXBatchedSequenceCache.h — the byte layer: map<seq_id, LayerPools> built lazily on a sequence's first write, attend looping the spans and concatenating in declared order, seq_clone copying the donor's pools.

backends/mlx/runtime/MLXCellCache.h, MLXInterpreter.h, MLXBackend.cpp — cell gets attend over its update_and_fetch, exec_update_and_attend shrinks to one call, the third builder is registered.

backends/mlx/examples/llm/run_llm_batched.cpp — --cache_kind, defaulting to batched-cell.

extension/llm/cache/test/cache_test.cpp, backends/mlx/test/mlx_batched_sequence_cache_test.cpp, backends/mlx/test/CMakeLists.txt, mlx_cell_cache_test.cpp — neutral tests for spans, placement, forks, context and admission; op-level parity test whose oracle is an MLXSequenceCache fed the same tokens.

Testing

extension_llm_cache_test 51, mlx_batched_sequence_cache_test 9, mlx_cell_cache_test 11, mlx_sequence_cache_test 16, test_update_and_attend.py 59.

pytorch-bot · 2026-09-10T00:40:41Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

⏳ 2 Pending, 2 Unrelated Failures

As of commit 54d4915 with merge base 3faa2e4 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Test QNN Windows MSVC build / build-qnn-windows-arm64 (gh) (trunk failure)
Test QNN Windows MSVC build / build-qnn-windows-x64 / windows-job (gh) (trunk failure)

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

github-actions · 2026-09-10T00:41:26Z

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-09-10T17:36:00Z


-// Lifecycle and admission, tensor-free.
+// Lifecycle, tensor-free. Admission needs a subject, so it belongs to the
+// faces below, which have one.


This comment is narrating why can_extend moved out of cache, but comments should not document history in this way. Keep them clear and concise and code self-documenting.

metascroy · 2026-09-10T17:37:27Z

+  virtual bool can_extend(int n = 1) const = 0;
  // Truncate to new_len; false = cannot grow, or the target is older than an
  // evicting layer still retains.
  virtual bool rewind(int new_len) = 0;


Maybe call position? It's the same thing conceptually

metascroy · 2026-09-10T17:38:53Z

+  // the model was trained to. Sequences are bounded independently, so this
+  // holds per sequence across a whole step.
+  // False also = an unknown sequence.
+  virtual bool can_admit(int32_t seq_id, int n = 1) const = 0;


Should this have same signature as declare_step? So user doesn't have to translate between the two formats?

metascroy · 2026-09-10T17:40:58Z

+  virtual bool seq_rm(int32_t seq_id) = 0;
+  // Truncate one sequence to new_len. False = an unknown sequence, a target
+  // that would grow it, or one older than an evicting layer still retains.
+  virtual bool rewind(int32_t seq_id, int new_len) = 0;


maybe call as int position? Same thing conceptually

metascroy · 2026-09-10T17:42:31Z

+  // Truncate one sequence to new_len. False = an unknown sequence, a target
+  // that would grow it, or one older than an evicting layer still retains.
+  virtual bool rewind(int32_t seq_id, int new_len) = 0;
  virtual int seq_len(int32_t seq_id) const = 0; // slots the sequence owns


Do we need both seq_len/next_pos?

Could we just have pos(int32_t seq_id) that reports the current position? This ties nicely with signature for rewind

Removed seq_len and renamed next_pos

metascroy · 2026-09-10T17:51:36Z

+  // Positions any one sequence may reach, from the model rather than from the
+  // caller: past it the model was never trained. Bounds a sequence, where
+  // capacity bounds them all together. Unset = only capacity bounds a sequence.
+  std::optional<int> max_context;


Thinking more here, maybe cache shouldn't be in business of max_context and can_admit/can_extend should be purely capacity based?

I agree, since can_admit was just checking max context, I removed that as well.

metascroy · 2026-09-10T17:52:56Z

+  // A fork: same capacity and policies, and the source's history. The byte
+  // layer copies the cells. Policies are rebuilt rather than shared, so the
+  // two sequences cannot disturb each other.
+  SequenceCache(const SequenceCache& other)


I suppose this is to support clone? If so, shouldn't we have an optional second arg for upto?

Yes, it's for seq_clone, we copy the full history and then rewind, so it can be rejected for evicted parts.

metascroy · 2026-09-10T17:59:52Z

+    }
+    auto it = rows_.find(src);
+    if (it != rows_.end()) {
+      rows_.emplace(*dst, it->second);


upto is not being applied correctly

metascroy · 2026-09-10T18:01:50Z

+// Apply `spec` to `q`: match the stored K/V to the query dtype, translate the
+// mask kind into what MLX takes, and run SDPA.
+inline Tensor
+attend(const AttendSpec& spec, const Tensor& q, float scale, StreamOrDevice s) {


what is the point of AttendSpec here? My not just pass args directl?

I think update_and_fetch is still useful to plan the runs, dtype and mask switch

metascroy · 2026-09-10T18:07:01Z

+    int at = 0;
+    for (const cache::SeqSpan& span : *spans) {
+      LayerPools& row = row_for(span.seq_id);
+      write_runs(


Is this correct for both sliding window and full linear attention?

Could we reuse the battle-tested logic in MLXSequenceCache.h by just having batched sequence cache hold std::map<seq_id, MLXSequenceCache>, then just directly calling attend in the for loop?

metascroy · 2026-09-10T18:07:59Z

    "KV storage dtype: bf16, fp16, or fp32");
+DEFINE_string(
+    cache_kind,
+    ::executorch::extension::llm::cache::kind::kBatchedCell,


why is batchedcell default when you said batchedsequence was faster?

metascroy · 2026-09-10T18:22:13Z

In batched_sequence_cache.cpp, BatchControl::declare_step()’s contract that a rejected declaration changes nothing. After a valid declaration, attempting one with an unknown sequence ID clears spans_ and declared_, so the previously valid step can no longer be placed. Build and validate candidate spans first, then replace the active declaration only on success.

This could be really simple if can_admit and declare_step had the same signature. Then declare_step could rely on can_admit for validation by calling it first thing.

metascroy · 2026-09-10T18:22:19Z

@claude review this code

claude · 2026-09-10T18:23:13Z

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

Code Review

I read the new control faces (cache.h), the batched-sequence layout (batched_sequence_cache.{h,cpp}), the MLX byte layers, and the cell layout for parity. Overall this is careful, well-documented work — the face machinery and the flat handoff are clean. A few substantive issues, in priority order.

1. `declare_step` is not atomic — a rejected step destroys the prior valid one (confirmed bug)

batched_sequence_cache.cpp:34-68 calls invalidate_step() as its first line, then validates. So a declaration that later fails (unknown id, no room, out of context) has already wiped a previously valid, un-placed declaration. This contradicts the contract you wrote in cache.h:110:

Also the admission gate: false = rejected and nothing changed

Build and validate the candidate spans in locals first, then swap them in only on success:

bool BatchedSequenceCache::declare_step(const std::vector<int32_t>& seq_ids) {
  if (seq_ids.empty() || !has_room(static_cast<int>(seq_ids.size()))) {
    return false;
  }
  std::map<int32_t, int> taking;
  for (int32_t id : seq_ids) {
    if (rows_.find(id) == rows_.end()) return false;
    ++taking[id];
  }
  for (const auto& [id, n] : taking) {
    if (!within_context(id, n)) return false;
  }
  // only now mutate:
  invalidate_step();
  // ... build spans_ ...
  declared_ = true;
  return true;
}

Note the cell layout has the same shape at cell_cache.cpp:59-82 — it mutates step_seq_ids_, declared_, and calls invalidate_steps() after passing its checks, but the checks come first there so it's safe. Worth a regression test asserting a good declaration survives a subsequent bad one on both layouts. Fix this →

Metascroy's follow-on suggestion (unify can_admit and declare_step signatures so the latter validates via the former) is worth doing, but note they aren't currently equivalent gates: can_admit is per-sequence context only and never consults has_room (shared pool capacity). If you route declare_step through can_admit, keep the pool check separate — which is what the comment at cache.h:118-119 already documents as intended.

2. `seq_clone` + `upto` ignores the prefix in the MLX byte layer (confirmed by reviewer)

The neutral layer correctly truncates the fork to len = min(upto, length) (batched_sequence_cache.cpp:92-100), but MLXBatchedSequenceCache::clone_bytes (MLXBatchedSequenceCache.h:122-124) copies the whole donor pools:

void clone_bytes(int32_t src, int32_t dst) override {
  rows_.emplace(dst, rows_.at(src));
}

For a flat layer this is harmless — the fork reads only [0, len) and later writes overwrite from len on. But for a ring layer it is not obviously correct: the ring addresses slots by position % ring_size, and when the donor has wrapped, the physical slot for a position p in the fork's needed window [len-window+1, len) may have been overwritten by a later donor position p + ring_size. The rewind(len) guard only ensures len >= length - window, which does not rule this out (e.g. small max_write). The comment at MLXBatchedSequenceCache.h:118-121 asserts whole-pool copy is fine "because a ring addresses slots by position" — but that reasoning is exactly what breaks once the ring has wrapped past the fork point.

There's no test exercising a partial (upto < length) clone against a ring layer at the byte level — mlx_batched_sequence_cache_test.cpp should add one. Either prove correctness with a test or restrict upto-clones to flat layers / narrow the copy.

3. Reviewer's architectural question on `MLXBatchedSequenceCache` (worth a serious look)

Metascroy asked whether the byte layer could hold std::map<seq_id, MLXSequenceCache> and just call the battle-tested attend() in the span loop, rather than reimplementing write/read/spec over raw LayerPools. Right now attend() (MLXBatchedSequenceCache.h:75-98) hand-rolls write_runs/read_runs/spec_for per span — which is exactly MLXSequenceCache::update_and_fetch's body. Delegating would:

eliminate the duplicated write/read/mask logic,
make issue Re-sync with internal repository #2 disappear (each per-seq MLXSequenceCache handles its own ring/window/clone),
and let clone_bytes become a MLXSequenceCache copy, which the copy-ctor at sequence_cache.h:165 already supports.

The cost is one SequenceCache of neutral bookkeeping per sequence (duplicating what rows_ in the base already tracks). Given #2, I'd lean toward this refactor.

Minor / style (aligning with your inline comments)

cache.h:78 — the comment narrating why can_extend moved out of Cache documents history; per CLAUDE.md ("Comments only for non-obvious global context") drop it.
cache.h:138-140 — seq_len vs next_pos: for this layout they're identical (next_pos just returns seq_len, batched_sequence_cache.cpp:134-136). They diverge only for the cell layout. A single pos(seq_id) would be simpler and, as you noted, pairs naturally with rewind(seq_id, new_len) — but it's a cross-layout API change, so fine to defer.
MLXCache.h:29-34 (AttendSpec) — reviewer asked why the struct vs. passing args directly. The struct earns its keep: spec_for/update_and_fetch build it in one place and the free attend() consumes it, so the dtype-match + mask-mode switch live once. I'd keep it.
run_llm_batched.cpp:48-52 — default is kBatchedCell while the PR summary says batched-sequence wins at B≥8. Reviewer flagged this; either default to the faster layout for the batched runner's typical B, or add a one-line justification (crossover / safety of the incumbent).
MLXBatchedSequenceCache.h:16 / header comment says "no mask for a decode token, causal for a prefill, and only a window narrower than the span needs one." Since spans are per-sequence and each is a contiguous run, that holds — good. Just confirm the ring/window path is actually reached by a test in mlx_batched_sequence_cache_test.cpp (the parity test's oracle is MLXSequenceCache, so a windowed multi-token span per sequence would exercise it).

Nothing here blocks the direction — #1 is a straightforward correctness fix and #2/#3 are the ones I'd resolve before landing.
· branch mlx-batched-seq-cache

metascroy · 2026-09-11T16:32:09Z

+      return std::nullopt;
+    }
+
+    MLXSequenceCache fork(it->second);


Should the constructor itself take upto? THis is unecessarily copying the whole history and then rewinding, but we only need to copy the slice

metascroy · 2026-09-11T17:06:15Z

In MLXBatchedSequenceCache.h:94–99, fork.rewind(len) lowers the logical length without preserving the physical retention floor. With window=4 and max_write=1, decode A to length 10, clone A at length 9 into B, then clone B at length 8 into C. Both clones succeed, but the copied ring contains only positions 6–9. C’s next attention at position 8 needs positions 5–8 and instead reads position 9’s KV from position 5’s slot. Preserve the actual retained-history bound across copies and rewinds so the second fork is rejected.

metascroy · 2026-09-11T20:12:10Z

+    if (upto <= 0 || !rewind(upto)) {
+      throw std::runtime_error("fork: position is not one this can rewind to");
+    }
+    kpool_.reserve(other.kpool_.size());


should we be reserving other.size?

something like max(initial_capacity, upto)?

This is actually the vector's element count: one Pool per layer, so n_layers. The slot sizing is separate, in clone_prefix: max(upto, initial_slots)

kiymetakdemir requested review from digantdesai and metascroy September 10, 2026 00:40

kiymetakdemir requested review from kirklandsign, larryliu0820 and mergennachin as code owners September 10, 2026 00:40

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 Sep 10, 2026

kiymetakdemir force-pushed the mlx-batched-seq-cache branch from 91cac31 to cf88ee4 Compare September 10, 2026 17:12

metascroy reviewed Sep 10, 2026

View reviewed changes

kiymetakdemir added 2 commits September 10, 2026 17:43

Add a batched-sequence KV cache layout and reorganize the control faces

cc4e79e

Make the byte copy a hook

bbfa436

kiymetakdemir force-pushed the mlx-batched-seq-cache branch 2 times, most recently from 4a6adfd to b7dbde1 Compare September 11, 2026 01:08

Hold an MLXSequenceCache per row, and drop max_context from the cache

ed94d84

kiymetakdemir force-pushed the mlx-batched-seq-cache branch from b7dbde1 to ed94d84 Compare September 11, 2026 01:29

metascroy reviewed Sep 11, 2026

View reviewed changes

kiymetakdemir added 3 commits September 11, 2026 10:40

Add a batched cache kind for the default batch layout

2c94722

Measure the rewind floor from what was written, not the length

e41ef87

Give a partial fork a pool holding just the prefix

54d4915

kiymetakdemir force-pushed the mlx-batched-seq-cache branch from b8660ce to 54d4915 Compare September 11, 2026 20:05

metascroy reviewed Sep 11, 2026

View reviewed changes

metascroy approved these changes Sep 11, 2026

View reviewed changes

kiymetakdemir merged commit 41fbae3 into pytorch:main Sep 11, 2026
255 of 260 checks passed

kiymetakdemir deleted the mlx-batched-seq-cache branch September 11, 2026 22:10

Conversation

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

Uh oh!

Uh oh!

pytorch-bot Bot commented Sep 10, 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/22662

⏳ 2 Pending, 2 Unrelated Failures

Uh oh!

github-actions Bot commented Sep 10, 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!

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!

Choose a reason for hiding this comment

Uh oh!

metascroy commented Sep 10, 2026

Uh oh!

metascroy commented Sep 10, 2026

Uh oh!

claude Bot commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review

1. declare_step is not atomic — a rejected step destroys the prior valid one (confirmed bug)

2. seq_clone + upto ignores the prefix in the MLX byte layer (confirmed by reviewer)

3. Reviewer's architectural question on MLXBatchedSequenceCache (worth a serious look)

Minor / style (aligning with your inline comments)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metascroy commented Sep 11, 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!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kiymetakdemir commented Sep 10, 2026 •

edited

Loading

pytorch-bot Bot commented Sep 10, 2026 •

edited

Loading

This PR needs a `release notes:` label

claude Bot commented Sep 10, 2026 •

edited

Loading

1. `declare_step` is not atomic — a rejected step destroys the prior valid one (confirmed bug)

2. `seq_clone` + `upto` ignores the prefix in the MLX byte layer (confirmed by reviewer)

3. Reviewer's architectural question on `MLXBatchedSequenceCache` (worth a serious look)

metascroy commented Sep 11, 2026 •

edited

Loading