Add functional update_and_attend KV-cache op + reference cache by kiymetakdemir · Pull Request #21112 · pytorch/executorch

kiymetakdemir · 2026-07-22T00:03:38Z

Summary
Adds a neutral kvcache::update_and_attend custom op for KV-cache attention. The op is functional to the tracer with the cache held off-graph via a process-global registry, so the exported .pte does not bake in cache size etc. Ships an eager reference cache (ContiguousReferenceCache) with runtime-selectable static/dynamic sizing and a hard capacity bound.

op_update_and_attend.py (new) — the kvcache::update_and_attend custom op and the off-graph cache registry
(install_cache/set_active/uninstall_cache).
op_update_and_attend_reference.py (new) — CacheConfig/AttendSpec and ContiguousReferenceCache (update_and_fetch, static/dynamic sizing, hard capacity) plus the neutral attend SDPA mechanism.
test_update_and_attend.py (new) — functional-graph + prefill/decode parity tests across static and dynamic sizing.
targets.bzl (modified) — adds the update_and_attend_py python_library.
BUCK (modified) — adds the test_update_and_attend python_test.

Test plan
pytest -v extension/llm/custom_ops/test_update_and_attend.py
Verifies the exported graph is functional (no buffer inputs/mutations, one op call per layer) and that outputs match cacheless causal attention for prefill and incremental decode under both static and dynamic sizing.

pytorch-bot · 2026-07-22T00:03:42Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Unrelated MacOS MPS test jobs time out (30 min) due to stale cpp_extension lock file poisoning self-hosted runners

✅ No Failures

As of commit 265ed7d with merge base 3ba685e ():
💚 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-22T00:04:29Z

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-22T17:21:06Z

+    ],
+)
+
+fbcode_target(_kind = runtime.python_binary,


Looks like you need to rebase?

Fixed this.

metascroy · 2026-07-22T17:23:57Z

+from executorch.extension.llm.custom_ops.op_update_and_attend_reference import attend
+
+_REGISTRY: Dict[str, object] = {}
+_CURRENT_KEY: Optional[str] = None


I know this is just a python reference, but maybe wrap this in a CacheRegistry class for neatness?

metascroy · 2026-07-22T17:26:19Z

+    scale: float,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    return torch.empty_like(q, dtype=out_dtype)


Is this the right shape? Shouldn't final dim be based on v? See https://docs.pytorch.org/docs/2.13/generated/torch.nn.functional.scaled_dot_product_attention.html

You are right, attn_weight is [B, H, q_len, k_len] and value is [B, H, k_len, D_v], so attn_weight @ value is [B, H, q_len, D_v].

metascroy · 2026-07-22T17:26:59Z

+
+
+@torch.library.custom_op("kvcache::update_and_attend", mutates_args=())
+def update_and_attend(


Can we have a doc string describing shape of each arg?

For example, the op should natively handle GQA, so make sure q/k/v shapes are described as such

metascroy · 2026-07-22T17:28:49Z

+        k = k.repeat_interleave(rep, dim=1)
+        v = v.repeat_interleave(rep, dim=1)
+
+    qf = q.to(torch.float32)


Why not use F.scaled_dot_product_attention?

metascroy · 2026-07-22T17:36:39Z

+    def _export(self, seq_len):
+        x = torch.randn(1, seq_len, self.hidden)
+        pos = _positions(0, seq_len)
+        idx = torch.arange(seq_len, dtype=torch.long)


Can we run functionalization on the graph? Sometimes export succeeds where functionalization fails (as an aside, ET always runs functionalization):

ep = ep.run_decompositions({})

metascroy · 2026-07-22T17:37:34Z

+            with self.subTest(sizing=sizing):
+                cache = ContiguousReferenceCache(self._config(sizing, cap))
+                install_cache(self.cache_key, cache)
+                set_active(self.cache_key)


What is difference between set_active and install?

install registers a cache under a key once (at creation), while active selects which installed cache the op uses for a given forward. They look identical with one cache, but distinct with multi-sequence, where we install each once and switch active per step.

metascroy · 2026-07-22T17:38:29Z

@@ -0,0 +1,224 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.


Make sure this is run by the unittest job in CI (look at raw logs to confirm)

metascroy · 2026-07-22T17:39:56Z

@@ -0,0 +1,167 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.


I take it this is just an example implementation?

At runtime we do not have to use things like AttendSpec for implementation?

Correct, and for the contiguous cases here the mask is always causal so it is not really functional now, I put it to use later for explicit but I can drop it for simplicity.

…che and tests

…nctionalization test

metascroy · 2026-07-22T21:24:22Z

Looks good! Thinking ahead to where c++ stuff will live, maybe it's better to move these files to a dedicated export/llm/cache folder.

Also mark as experimental

kiymetakdemir requested review from SS-JIA, digantdesai and metascroy July 22, 2026 00:03

kiymetakdemir requested review from larryliu0820 and mergennachin as code owners July 22, 2026 00:03

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 22, 2026

digantdesai reviewed Jul 22, 2026

View reviewed changes

Comment thread extension/llm/custom_ops/op_update_and_attend_reference.py Outdated

metascroy reviewed Jul 22, 2026

View reviewed changes

Comment thread extension/llm/cache/test_update_and_attend.py

metascroy reviewed Jul 22, 2026

View reviewed changes

Add functional off-graph update_and_attend KV-cache op + reference ca…

8c72c6b

…che and tests

kiymetakdemir force-pushed the update-and-attend-op-phase1 branch from c47a9f3 to 8c72c6b Compare July 22, 2026 17:52

kiymetakdemir and others added 2 commits July 22, 2026 13:11

Address review: CacheRegistry, F.sdpa, fake v-dim fix, docstrings, fu…

95e9df5

…nctionalization test

Merge branch 'main' into update-and-attend-op-phase1

6a2a5f9

metascroy approved these changes Jul 22, 2026

View reviewed changes

kiymetakdemir and others added 2 commits July 22, 2026 14:43

Move update_and_attend cache to extension/llm/cache; mark experimental

ff56fed

Merge branch 'main' into update-and-attend-op-phase1

265ed7d

kiymetakdemir merged commit 3389d51 into pytorch:main Jul 23, 2026
297 of 298 checks passed

kiymetakdemir deleted the update-and-attend-op-phase1 branch July 23, 2026 17:23



		@torch.library.custom_op("kvcache::update_and_attend", mutates_args=())
		def update_and_attend(

		@@ -0,0 +1,224 @@
		# Copyright (c) Meta Platforms, Inc. and affiliates.

		@@ -0,0 +1,167 @@
		# Copyright (c) Meta Platforms, Inc. and affiliates.

Uh oh!

Conversation

kiymetakdemir commented Jul 22, 2026

Uh oh!

pytorch-bot Bot commented Jul 22, 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/21112

❗ 1 Active SEVs

✅ No Failures

Uh oh!

github-actions Bot commented Jul 22, 2026

This PR needs a release notes: label

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

metascroy Jul 22, 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!

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 Jul 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

pytorch-bot Bot commented Jul 22, 2026 •

edited

Loading

This PR needs a `release notes:` label

metascroy Jul 22, 2026 •

edited

Loading