fix: stop get_cache handing out the stored KV cache by reference - #2303
fix: stop get_cache handing out the stored KV cache by reference#2303jeojdi1 wants to merge 2 commits into
Conversation
`_concat_caches` returns `caches[0]` unchanged when a single cache id is requested, so `get_cache` hands the caller the stored `KVCacheItem.memory` object itself. The caller passes that cache to `generate`, which appends to it in place, so the saved activation memory grows on every chat turn -- measured stored length 6 -> 19 -> 32 -> 45 over three turns. Nothing in the API suggests that retrieving a memory mutates it, and the multi-cache path already builds a fresh container, so only this early return leaked the reference. Returns a new `DynamicCache` sharing the same tensors instead. The tensors are not cloned: `generate` appends along the sequence axis rather than writing into existing rows, so a fresh container is enough to protect the stored item without paying to duplicate the cache. Adds `test_get_cache_does_not_alias_stored_memory`, which simulates `generate` by appending to the returned cache and asserts the stored length is unchanged. It fails on the current code and passes with this change.
Memtensor-AI
added
area:memory
🤖 Open Code ReviewTarget: PR #2303 🔍 OpenCodeReview found 1 issue(s) in this PR. 1.
|
|
Review feedback: _copy_cache performs only Python-level attribute copies, so the local torch import was never referenced.
|
Thanks for the review — one accepted, one I'd push back on. 1. Unused 2. The conditional
def lazy_initialization(self, key_states, value_states) -> None:
self.dtype, self.device = key_states.dtype, key_states.device
self.keys = torch.tensor([], dtype=self.dtype, device=self.device)
self.values = torch.tensor([], dtype=self.dtype, device=self.device)
self.is_initialized = TrueThe load-bearing line is the last one. if not self.is_initialized or self.keys.numel() == 0:
return 0So a layer that gets On the second option — an I've left the ordering as-is, but it clearly isn't obvious from reading it — happy to add a comment making the dependency explicit if you'd like. |
|
5 participants
Description
_concat_cachesreturnscaches[0]unchanged when a single cache id is requested, soget_cache()hands the caller the storedKVCacheItem.memoryobject itself. The caller passes that cache togenerate(), which appends to it in place — so the saved activation memory grows on every chat turn.Measured across three turns, the stored cache length went 6 → 19 → 32 → 45. Nothing in the API suggests that retrieving a memory mutates it, and the multi-cache path already builds a fresh container, so only this early return leaks the reference.
This PR returns a new
DynamicCachesharing the same tensors. The tensors are deliberately not cloned:generate()appends along the sequence axis rather than writing into existing rows, so a fresh container is sufficient to protect the stored item without paying to duplicate the cache.Related Issue (Required): #2301
Type of change
How Has This Been Tested?
test_get_cache_does_not_alias_stored_memoryasserts the returned object is not the stored one, then simulatesgenerate()by appending a step to the returned cache and asserts the stored length is unchanged.Verified fails-before / passes-after on
transformers 5.16.1:Checklist
A note on the target branch
CONTRIBUTING.mdsays to open PRs againstdev, but nodevbranch exists — onlymainanddev-v2.0.28…dev-v2.0.32. This is againstmain(185ebdb, "Dev v2.0.32"). Happy to retarget.