diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..ec95d8778 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -80,6 +80,40 @@ 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. + """ + 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 +240,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" + )