SIGN IN SIGN UP

Add functional update_and_attend KV-cache op (#21112)

**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.
K
Kiymet Akdemir committed
3389d51918773afb9fabb848c9b52b4bff042dd6
Parent: 2cce41f
Committed by GitHub <noreply@github.com> on 7/23/2026, 2:46:24 PM