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