From cd8a796741520f2420638326a21c8a04c823fc64 Mon Sep 17 00:00:00 2001 From: jeojdi1 Date: Thu, 27 Aug 2026 18:31:38 -0400 Subject: [PATCH 1/2] fix: stop get_cache handing out the stored KV cache by reference `_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. --- src/memos/memories/activation/kv.py | 44 ++++++++++++++++++++++++++- tests/memories/activation/test_kv.py | 45 ++++++++++++++++++++++++---- 2 files changed, 82 insertions(+), 7 deletions(-) diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..85edbeda4 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -80,6 +80,42 @@ def get_cache(self, cache_ids: list[str]) -> DynamicCache | None: return self._concat_caches(caches_to_merge) + @staticmethod + def _copy_cache(cache: DynamicCache) -> DynamicCache: + """Shallow-structure copy of a DynamicCache that shares no mutable state. + + The tensors themselves are not cloned -- ``generate`` appends along the + sequence axis rather than writing into the existing rows, so a fresh + container with the same tensor objects is enough to protect the stored + item, without paying to duplicate the cache. + """ + import torch # noqa: F401 (kept local, matching this module's style) + + copy = DynamicCache() + if hasattr(cache, "layers"): + if not hasattr(copy, "layers"): + copy.layers = [] + layer_cls = type(cache.layers[0]) if cache.layers else None + while layer_cls is not None and len(copy.layers) < len(cache.layers): + copy.layers.append(layer_cls()) + for i, src in enumerate(cache.layers): + dst = copy.layers[i] + if not getattr(dst, "is_initialized", True) and hasattr( + dst, "lazy_initialization" + ): + dst.lazy_initialization(src.keys, src.values) + dst.keys = src.keys + dst.values = src.values + elif hasattr(cache, "key_cache"): + for i in range(len(cache.key_cache)): + copy.key_cache.append(cache.key_cache[i]) + copy.value_cache.append(cache.value_cache[i]) + else: + raise AttributeError( + "DynamicCache object has neither 'layers' nor 'key_cache' attributes" + ) + return copy + def get(self, memory_id: str) -> KVCacheItem | None: """Get a memory by its ID. @@ -206,7 +242,13 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache: assert caches, "Need at least one cache" if len(caches) == 1: - return caches[0] + # Do NOT hand back the stored object. The caller passes this cache to + # ``generate``, which appends to it in place, so returning the stored + # item makes every chat turn grow the saved activation memory: + # observed stored length 6 -> 19 -> 32 -> 45 over three turns. The + # multi-cache path below already builds a new container, so only this + # early return leaked the reference. + return self._copy_cache(caches[0]) merged = DynamicCache() diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 6490d687f..6f5d1a35b 100644 --- a/tests/memories/activation/test_kv.py +++ b/tests/memories/activation/test_kv.py @@ -33,11 +33,11 @@ def kv_memory(dummy_config): yield KVCacheMemory(dummy_config) -def make_filled_cache(): - # Create a DynamicCache with at least one dummy tensor layer +def make_filled_cache(seq_len: int = 3, n_layers: int = 1): + """Create a DynamicCache with dummy tensors, on any transformers version.""" cache = DynamicCache() - cache.key_cache.append(torch.zeros(1, 2, 3)) - cache.value_cache.append(torch.zeros(1, 2, 3)) + for layer_idx in range(n_layers): + cache.update(torch.zeros(1, 2, seq_len, 4), torch.zeros(1, 2, seq_len, 4), layer_idx) return cache @@ -59,8 +59,11 @@ def test_get_cache_merge(kv_memory): merged = kv_memory.get_cache([item1.id, item2.id]) assert isinstance(merged, DynamicCache) # Check the number of layers in merged key/value cache - assert len(merged.key_cache) == 1 - assert len(merged.value_cache) == 1 + if hasattr(merged, "layers"): + assert len(merged.layers) == 1 + else: + assert len(merged.key_cache) == 1 + assert len(merged.value_cache) == 1 def test_delete_and_get_all(kv_memory): @@ -84,3 +87,33 @@ class DummyTextualMemory: item = kv_memory.from_textual_memory(DummyTextualMemory()) assert isinstance(item, KVCacheItem) assert item.metadata["bar"] == 1 + + +def test_get_cache_does_not_alias_stored_memory(kv_memory): + """A cache handed to the caller must not be the stored object itself. + + ``get_cache`` returned ``caches[0]`` directly whenever a single id was + requested. The caller passes that cache to ``generate``, which appends to it + in place, so the saved activation memory grew on every chat turn -- observed + as a stored length of 6 -> 19 -> 32 -> 45 across three turns. Nothing in the + API suggests retrieving a memory mutates it. + + Simulating ``generate`` by appending to the returned cache must leave the + stored item unchanged. Fails before this change, passes after. + """ + item = KVCacheItem(memory=make_filled_cache(seq_len=5, n_layers=2)) + kv_memory.add([item]) + stored_len_before = kv_memory.get(item.id).memory.get_seq_length() + + handed_out = kv_memory.get_cache([item.id]) + assert handed_out is not item.memory, "get_cache returned the stored object itself" + + # what generate() does: append one step of new keys/values + for layer_idx in range(2): + handed_out.update(torch.zeros(1, 2, 1, 4), torch.zeros(1, 2, 1, 4), layer_idx) + + stored_len_after = kv_memory.get(item.id).memory.get_seq_length() + assert stored_len_after == stored_len_before, ( + f"stored memory grew {stored_len_before} -> {stored_len_after} because the " + f"caller mutated the cache it was handed" + ) From a8b1bbf0785edaa6bb2d6c993293987b3117b825 Mon Sep 17 00:00:00 2001 From: jeojdi1 Date: Tue, 1 Sep 2026 04:18:31 -0400 Subject: [PATCH 2/2] refactor: drop unused torch import from _copy_cache Review feedback: _copy_cache performs only Python-level attribute copies, so the local torch import was never referenced. --- src/memos/memories/activation/kv.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 85edbeda4..ec95d8778 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -89,8 +89,6 @@ def _copy_cache(cache: DynamicCache) -> DynamicCache: container with the same tensor objects is enough to protect the stored item, without paying to duplicate the cache. """ - import torch # noqa: F401 (kept local, matching this module's style) - copy = DynamicCache() if hasattr(cache, "layers"): if not hasattr(copy, "layers"):