Add off-graph KV-cache cell reference: multi-sequence batching by kiymetakdemir · Pull Request #21761 · pytorch/executorch

kiymetakdemir · 2026-08-11T22:12:52Z

Summary

Adds CellReferenceCache to the eager KV-cache reference: many sequences over one pool of per-token cells. A cell holds one token's K/V plus that token's position and the set of sequences owning it, so a sequence needn't occupy a contiguous range and two sequences can share cells — a fork sets a second owner bit instead of copying K/V. The batch is flat on the token axis (B = 1).

Files

extension/llm/cache/reference_cache.py — CellReferenceCache beside ContiguousReferenceCache, a separate class rather than a flag. Per-cell position and owner bitset; a token claims the lowest free cell, which keeps the read window tight. Visibility is a property of the cell, not the layout — query i attends cell j iff j is occupied, shares a sequence with i, and is no newer than i — so the mask is a broadcast compare over the cell metadata rather than a host loop, and the spec is always EXPLICIT. begin_step declares the per-token sequences and decides admission before the forward; since cells are interchangeable, a step that passes that check cannot then fail to allocate. The first layer of a step allocates and memoizes the cells and mask and the rest reuse them, so allocation happens exactly once and there is no commit; a spent declaration plus a served-layer set turn a skipped begin_step into an error rather than a silent reuse of the previous step's sequences. seq_cp/seq_rm/seq_keep are bitset operations and a cell frees only once its owner set empties. DYNAMIC sizing grows the pool by appending rows, so a cell keeps the index the plan holds. Multi-dimensional positions are rejected rather than guessed at.
extension/llm/cache/test_update_and_attend.py — ten cases.

Testing

Ten cases added to the existing unittest suite:

a single sequence matches the cacheless baseline, so the degenerate case is not special-cased
two sequences prefilled in one forward, each matching the baseline it would have produced alone
a batched decode step continuing both sequences
a fork sharing the trunk's cells without consuming any, with the branch attending the trunk's history
seq_rm freeing only the cells no other sequence still owns
admission failing in begin_step, before the forward
the step protocol: a forward that skips begin_step, a declared/actual token count mismatch, and a sequence verb not disguising
a skipped declaration
growth keeping every cell's index and bytes, and STATIC agreeing with DYNAMIC

pytest extension/llm/cache/test_update_and_attend.py -q

pytorch-bot · 2026-08-11T22:12:56Z

🔗 Helpful Links

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

📄 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 490bfea with merge base e60faa2 ():
💚 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-08-11T22:13:44Z

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-08-11T22:23:58Z

@claude review this code

claude · 2026-08-11T22:24:47Z

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

Code review: off-graph KV-cache cell reference

Read reference_cache.py
Read the test file and update_and_attend.py
Analyze masking, allocation, sequence verbs, growth, step protocol
Post review

Nicely done — the design is coherent and the docstrings carry the intent well. I traced the masking, allocation, growth, and step-protocol paths and believe the core logic is correct. The mask predicate (occupied & shares-a-sequence & pos ≤ query-pos) cleanly subsumes causality, cross-sequence isolation, and forking; _claim → _ensure → scatter keeps cell indices stable across DYNAMIC growth; and the _declared/_served/_plan triad enforces the one-declaration-one-allocation protocol as the tests claim. AttendExplicitTest pinning the bool true = attend polarity is exactly right given SDPA's mask convention.

A few observations, none blocking:

1. Sequence verbs skip the range validation that begin_step enforces
begin_step guards 0 <= seq < MAX_SEQS, but seq_cp/seq_rm/seq_keep don't. An out-of-range dst (e.g. seq_cp(0, 70)) writes 1 << 70 into _seq[i] as a Python big-int; that stays invisible until _allocate does torch.tensor(self._seq[:n]), where it overflows int64 and surfaces as a confusing error far from the cause. Cheap to validate dst/seq at the verb, matching begin_step. reference_cache.py:270-297 Fix this →

2. A failed _allocate leaves the declaration unspent
When _allocate raises on a token-count mismatch or a multi-dim position (reference_cache.py:355-365), it returns before self._declared = False, so _declared stays True and no cells were claimed. That's memory-safe, but it means a second update_and_fetch on an unserved layer — with a now-matching position and no fresh begin_step — would succeed, which is slightly at odds with the strict "one declaration, one allocation" story the class otherwise tells. The existing test only checks the immediate raise. If the leniency is intentional, a one-line comment would help; if not, spending the declaration on any allocation attempt would make the protocol uniform.

3. Read window spans the whole occupied pool, not the per-sequence footprint
read_len = used_end, so a single-sequence decode still fetches and masks out every other sequence's cells. Correct and fine for a reference oracle, but worth flagging for the byte-layer implementation: mask width (and the SDPA cost) grows with the total pool rather than the attending sequence's history. The _claim full-capacity scan and free_cells/seq_len full scans are similar reference-only costs — already acknowledged in the comments.

4. seq_keep is unexercised
seq_cp/seq_rm each have a dedicated test, but seq_keep (the "commit one sequence, free the rest" verb) has none. Its &= bit logic looks right, but a small test (keep a forked sequence, assert the other's exclusive cells free while shared cells survive) would round out the suite and guard against a future &=→|= slip.

Minor / optional

Empty begin_step([]) is admitted (can_extend(0) is true) and yields an empty plan/mask. Harmless, but if a zero-token step is never meaningful, an early guard would make misuse obvious.

Everything above is polish on an otherwise solid change; the correctness-critical paths hold up.

metascroy · 2026-08-11T22:25:00Z

 mechanism (``attend`` below) is applied by the op/backend from that spec.

-Scope for this initial slice: single sequence, contiguous placement, float KV.
+Two caches share the op: ``ContiguousReferenceCache`` (one sequence appended in


It can be a follow-up PR, but didn't we rename ContiguousReferenceCache as SequenceCache or something in C++ code?

metascroy · 2026-08-11T23:09:06Z

+    def free_cells(self) -> int:
+        return self._pos.count(-1)
+
+    def can_extend(self, n: int = 1) -> bool:


What does can_extend mean in the context of multiple sequences?

Whether the total concatenated length is extendible?

Yes exactly, it checks if there are enough cells for this step's total tokens up to a capacity

metascroy · 2026-08-11T23:58:59Z

+        self._shrink()
+        self._invalidate_plan()
+
+    def seq_keep(self, keep: int) -> None:


Get rid of this for now. We can always extend API in future

metascroy · 2026-08-12T00:05:34Z

+    "update_and_attend KV cache is experimental and may change without notice."
+)
+class CellReferenceCache:
+    """Per-cell KV history for several sequences sharing one pool.


The API here is leaving a lot of bookkeeping to the caller to flatten sequences.

Can we have a utility that flattens {seq_id: (token_id, pos_id)} to the parallel arrays seq_id, token_id, pos_id?

metascroy · 2026-08-12T00:06:47Z

+
+    # -- internals ----------------------------------------------------------
+
+    def _allocate(self, position: torch.Tensor, device: torch.device) -> _CellStepPlan:


Why is device here?

It was there to build the plan's tensors on the same device as the K/V, since both are used against the pools. Now dropped it, it takes the device from the pools themselves (self._k[0].device).

metascroy · 2026-08-12T00:11:14Z

+                f"the forward carries {len(positions)}"
+            )
+        cells = [
+            self._claim(pos, 1 << seq) for pos, seq in zip(positions, self._step_seq)


How are prefixes being shared here?

Isn't 1 << seq just having one owners.

If (pos, seq, tok) all match, are we identify shared cell? How do we do prefix sharing with batch prefill?

Yes currently during write it is only one owner, and shared with seq_cp

metascroy · 2026-08-12T00:14:26Z

+        self._plan = None
+        self._served.clear()
+
+    def seq_cp(self, src: int, dst: int, p0: int = 0, p1: Optional[int] = None) -> None:


Are both p0/p1 definable if only prefix sharing is supported?

I dropped p0, I guess it has no use case.

metascroy · 2026-08-12T00:15:27Z

+    def can_extend(self, n: int = 1) -> bool:
+        return self.free_cells() >= n
+
+    def seq_len(self, seq: int) -> int:


There's inconsistent use of seq vs. seq_id in the code.

Let's consolidate on seq_id.

Done, also changed self._seq → self._owners

metascroy · 2026-08-12T00:16:24Z

+                self._seq[i] |= dst_bit
+        self._invalidate_plan()
+
+    def seq_rm(self, seq: int, p0: int = 0, p1: Optional[int] = None) -> None:


Document that None for p1 means until end?

metascroy · 2026-08-12T00:19:54Z

+        """Drop seq's claim; a cell frees only once no sequence owns it."""
+        bit = 1 << seq
+        for i in range(self._used_end):
+            if self._seq[i] & bit and self._in_range(self._pos[i], p0, p1):


If we start with [0, 10] populated for seq0 and seq1, sharing cell_ids [0, 10], and we then do seq_rm on seq1 from [2, 5], what does that mean for what remains?

Is that a valid operation?

That is valid, seq0 stays as it is, seq1 owns positions 0,1 and 5-10. One use case could be attention sink; keeping the first tokens and evicting the ones after sink.

metascroy · 2026-08-12T00:20:36Z

+            out, self.model.reference_forward(x, torch.arange(5)), atol=1e-4, rtol=1e-4
+        )
+
+    def test_batched_sequences_match_separate_runs(self):


I think these would be eaiser to read if we had the flatten utiltity mentioned above

metascroy · 2026-08-12T00:21:23Z

@claude review the code again, along with reviewer comments

claude · 2026-08-12T00:22:20Z

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

Re-review with reviewer comments

Read reference_cache.py and test file
Cross-check each of @metascroy's inline comments
Re-analyze correctness paths
Post consolidated review

Re-read the current revision and worked through each of @metascroy's threads. My take on where I agree, where I'd push, and a couple of things worth pinning down before this lands.

The one that matters most: prefix sharing vs. batch prefill (`:367`, `:270`, `:282`)

@metascroy's instinct is right, and I think it's worth being explicit in the design: this cache does not detect shared prefixes — sharing is only ever produced by seq_cp. _claim gives every token in a batch prefill its own lowest-free cell, and the mask predicate (cell_seq & tok_seq) != 0 only unifies cells that already carry both owner bits. So if you prefill seq0 and seq1 in one forward and they happen to share a prefix, you get two disjoint copies of that prefix, not one shared set. The only path to sharing is: prefill the trunk under one seq, then seq_cp(trunk, branch) to graft a second owner bit on — which is exactly what test_fork_shares_cells_and_history exercises. That's a coherent model (fork-after-prefill, à la vLLM's copy-on-fork), but the docstring's "two may share cells" reads as if the pool dedups automatically, and it doesn't. One sentence clarifying "sharing is created by seq_cp, never inferred" would remove the ambiguity the review thread surfaced.

That reframing also answers the :270/:282 questions:

p0/p1 on seq_cp (:270) — if the supported operation is prefix sharing, a prefix is [0, p1), so p0 is dead surface. I'd drop p0 (and let p1=None mean "whole sequence") rather than expose an arbitrary [p0, p1) window that nothing uses and that admits sharing a middle slice — a shape the rest of the design has no story for.
seq_rm from the middle (:282) — mechanically it's well-defined: clearing seq1's bit on positions [2,5) of a shared [0,10] just drops seq1's visibility of those cells (they stay resident because seq0 still owns them), leaving seq1 with a positional hole {0,1,5,6,…}. It's not incorrect — the mask is position-based so a hole attends fine — but it's a shape no realistic caller wants, and it exists only because the range API is more general than the use case. Narrowing the API (above) makes the question moot.

Trim the surface now, extend later

Remove seq_keep (:289) — agreed, and note this retires my earlier "seq_keep is untested" observation; deleting it is cleaner than adding a test for a verb you don't need yet.
Flatten helper (:192, test :313) — agreed this is worth adding. Every caller and every test hand-builds parallel positions/seq_ids arrays; a small flatten({seq_id: (token_ids, pos_ids)}) -> (token_ids, pos_ids, seq_ids) utility would both shrink the test bodies (test_batched_* especially) and make begin_step's contract self-evident. Keep it a plain function, not a method — it's caller-side bookkeeping.

Naming / docs (cheap, do them)

seq vs seq_id (:246) — consolidate on seq_id throughout (seq_len, seq_cp, seq_rm, the _allocate locals, MAX_SEQS loop var). Purely mechanical.
can_extend docstring (:243) — the thread resolved the meaning ("enough free cells for this step's total token count, up to capacity"); fold that one line into the docstring so the next reader doesn't have to re-derive it. It is deliberately sequence-agnostic — cells are interchangeable — which is the whole reason admission can be decided before positions arrive.
p1=None means "to end" (:278) — document on whichever verbs survive the trim.
device in _allocate (:350) — the honest answer for the thread: it exists so cells/mask land on the same device as k/v (passed as k.device), keeping the scatter and the SDPA mask off-host on GPU. It's legitimate, not vestigial — worth a half-line comment saying so, since everything else in the class is host-side Python lists and the lone device arg looks out of place without it.
ContiguousReferenceCache → match C++ SequenceCache (:23) — fine as the follow-up you suggested; renaming both references in one pass later is cleaner than a half-rename now.

Carrying forward from my first pass (still open, independent of the threads)

seq_cp/seq_rm don't range-check the seq id. begin_step guards 0 <= seq < MAX_SEQS; the verbs don't, so seq_cp(0, 70) writes 1 << 70 as a Python big-int and only blows up later in _allocate's torch.tensor(self._seq[:n]) int64 overflow — far from the cause. Add the same guard to whichever verbs remain. reference_cache.py:270-297
A failed _allocate leaves the declaration unspent. The token-count-mismatch and multi-dim-position raises (:355-365) return before self._declared = False (:380), so a retry on the same unserved layer with a corrected position succeeds without a fresh begin_step — slightly at odds with the "one declaration, one allocation" story. If intentional, a one-line comment; if not, spend the declaration on any allocation attempt. reference_cache.py:350-383
Empty begin_step([]) is admitted and yields an empty plan/mask. Harmless; an early guard would make a zero-token step read as misuse rather than silently succeed.

Correctness

I re-traced allocation, the mask predicate, DYNAMIC growth, cross-sequence isolation, and the _declared/_served/_plan protocol against the ten cases — the correctness-critical paths hold. Freed middle cells (_pos[i] == -1) are correctly excluded by cell_pos >= 0; index stability across growth holds; isolation between co-batched sequences holds. Everything above is API-shape and documentation polish, plus the two small robustness items (1) and (2). No behavioral bug found in the paths the tests cover.
• branch kvcache-cells

…verbs

digantdesai · 2026-08-12T19:07:04Z

+            self.model(tokens, positions, logits_indices)
+
+        tokens, positions, seq_ids, logits_indices = flatten_step(
+            {0: (a[:, 3:], 3), 1: (b[:, 2:], 2)}


who will keep track of these positions?

Right now the caller does. The cache could report it instead: next_pos(seq_id) = one past the highest position among the cells that sequence owns (0 if it owns none), so the call becomes flatten_step({0: (tok, cache.next_pos(0))}) and the caller stops tracking positions itself.

digantdesai · 2026-08-12T20:05:32Z

+                self._owners[i] &= ~bit
+                if self._owners[i] == 0:
+                    self._pos[i] = -1
+        self._shrink()


I guess this will wind it back id self._owners == 0 (when the last seq_id is removed)? if yes, do we have a test? if no, can we add? For memory leaks.

Yes _shrink walks used_end back past every trailing cell whose owners went empty. Now I added test_freeing_the_tail_shrinks_the_read_window

digantdesai · 2026-08-12T20:52:39Z

        )
        return q, k, v

    def forward(self, x, position, logits_indices):


can you remind me why position and logits_indices has to be in the forward sign?

position is for RoPE and logits_indices for lm head

digantdesai · 2026-08-12T21:31:17Z


-Scope for this initial slice: single sequence, contiguous placement, float KV.
+Two caches share the op: ``ContiguousReferenceCache`` (one sequence appended in
+place) and ``CellReferenceCache`` (many sequences over a pool of per-token cells,


did you consider cell size to be >1 token? This is in the context of the paged attn impls.

Add off-graph KV-cache cell reference: multi-sequence batching

a0c7e3d

kiymetakdemir requested review from JakeStevens, digantdesai and metascroy August 11, 2026 22:12

kiymetakdemir self-assigned this Aug 11, 2026

kiymetakdemir requested review from larryliu0820 and mergennachin as code owners August 11, 2026 22:12

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 Aug 11, 2026

metascroy reviewed Aug 11, 2026

View reviewed changes

metascroy reviewed Aug 12, 2026

View reviewed changes

Address review: flatten_step helper, seq_id guards, tighter sequence …

91c3901

…verbs

metascroy approved these changes Aug 12, 2026

View reviewed changes

digantdesai reviewed Aug 12, 2026

View reviewed changes

Address review: test that freeing the tail shrinks the read window

490bfea

kiymetakdemir force-pushed the kvcache-cells branch from 7e83627 to 490bfea Compare August 12, 2026 22:14

kiymetakdemir merged commit dee4fd4 into pytorch:main Aug 13, 2026
184 checks passed

kiymetakdemir deleted the kvcache-cells branch August 13, 2026 16:51


		# -- internals ----------------------------------------------------------

		def _allocate(self, position: torch.Tensor, device: torch.device) -> _CellStepPlan:

Uh oh!

Conversation

kiymetakdemir commented Aug 11, 2026

Uh oh!

pytorch-bot Bot commented Aug 11, 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/21761

✅ No Failures

Uh oh!

github-actions Bot commented Aug 11, 2026

This PR needs a release notes: label

Uh oh!

metascroy commented Aug 11, 2026

Uh oh!

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

Uh oh!

Code review: off-graph KV-cache cell reference

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metascroy commented Aug 12, 2026

Uh oh!

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

Uh oh!

Re-review with reviewer comments

The one that matters most: prefix sharing vs. batch prefill (:367, :270, :282)

Trim the surface now, extend later

Naming / docs (cheap, do them)

Carrying forward from my first pass (still open, independent of the threads)

Correctness

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

pytorch-bot Bot commented Aug 11, 2026 •

edited

Loading

This PR needs a `release notes:` label

claude Bot commented Aug 11, 2026 •

edited

Loading

claude Bot commented Aug 12, 2026 •

edited

Loading

The one that matters most: prefix sharing vs. batch prefill (`:367`, `:270`, `:282`)

digantdesai Aug 12, 2026 •

edited

Loading