fix(activation): keep the stored KV cache read-only during generation (#2301)

src/memos/llms/hf.py

@@ -82,7 +82,7 @@ def generate(
         if past_key_values is None:
             return self._generate_full(prompt, **kwargs)
         else:
-            return self._generate_with_cache(prompt, past_key_values, **kwargs)
+            return self._generate_with_cache(prompt, self._clone_cache(past_key_values), **kwargs)
 
     def generate_stream(
         self, messages: MessageList, past_key_values: DynamicCache | None = None, **kwargs
@@ -102,7 +102,47 @@ def generate_stream(
         if past_key_values is None:
             yield from self._generate_full_stream(prompt)
         else:
-            yield from self._generate_with_cache_stream(prompt, past_key_values)
+            yield from self._generate_with_cache_stream(prompt, self._clone_cache(past_key_values))
+
+    @staticmethod
+    def _clone_cache(cache: DynamicCache) -> DynamicCache:
+        """
+        Copy a KV cache before it is handed to the model.
+
+        ``transformers`` appends the generated key/value tensors to the cache
+        object in place, so the cache that is stored in the activation memory
+        would otherwise grow with every turn and the growth would be written
+        back to disk on the next dump.
+
+        The copy is the price of keeping the stored cache untouched: one extra
+        copy of the activation cache is alive while the model generates. That
+        is the same order as the prefill work itself, and it is preferred over
+        cropping the stored cache back to its original length afterwards,
+        which would mutate the stored cache during the generation and leave
+        it grown whenever the generation is interrupted.
+
+        A cache that is already an owned copy -- one that was handed out by
+        the activation memory and therefore shares no tensor with the stored
+        cache -- is not copied again: the copy would be pure overhead and
+        would keep a second full set of K/V tensors alive during the
+        generation. The stored cache stays read-only either way, because an
+        owned copy never aliases it.
+
+        Args:
+            cache (DynamicCache): The cache provided by the caller.
+
+        Returns:
+            DynamicCache: A cache that does not share any tensor with the
+            stored activation memory.
+        """
+        # Imported lazily: memos.llms.factory imports this module, so a
+        # module-level import of memos.memories.activation.kv would be circular.
+        from memos.memories.activation.kv import clone_dynamic_cache, is_owned_cache
+
+        if is_owned_cache(cache):
+            return cache
+
+        return clone_dynamic_cache(cache)
 
     def _generate_full(self, prompt: str, **kwargs) -> str:
         """

src/memos/memories/activation/kv.py

@@ -1,7 +1,9 @@
+import copy
 import os
 import pickle
 
 from datetime import datetime
+from typing import Any
 
 from transformers import DynamicCache
 
@@ -13,6 +15,50 @@
 from memos.memories.textual.item import TextualMemoryItem
 
 
+#: Attribute that marks a cache as an *owned copy*.
+#:
+#: A cache that leaves the activation memory is an owned copy when it shares
+#: no tensor with the stored cache: the caller may hand it to a model that
+#: appends to it in place without corrupting the stored cache. The marker
+#: records that ownership, so the model boundary can tell an already owned
+#: cache from a borrowed one and does not copy it a second time.
+_OWNED_CACHE_MARKER = "_memos_owned_cache_copy"
+
+
+def is_owned_cache(cache: Any) -> bool:
+    """
+    Whether ``cache`` is an owned copy that the caller may mutate in place.
+
+    A cache is owned when it was produced by :func:`clone_dynamic_cache` or by
+    the merge path of ``_concat_caches``: it shares no tensor with the
+    activation memory, so handing it to a model that appends to it in place
+    cannot corrupt the stored cache. Any other cache is *borrowed* and has to
+    be copied before it is given to the model.
+
+    Args:
+        cache: The cache to inspect.
+
+    Returns:
+        bool: True when the cache is an owned copy.
+    """
+    return bool(getattr(cache, _OWNED_CACHE_MARKER, False))
+
+
+def _mark_cache_as_stored(cache: Any) -> None:
+    """
+    Drop the ownership marker of a cache that is stored in the memory.
+
+    A stored cache belongs to the activation memory, not to the caller: it
+    must never be treated as an owned copy, otherwise the model would append
+    to it in place and grow the stored memory again.
+
+    Args:
+        cache: The cache that is about to be stored.
+    """
+    if getattr(cache, _OWNED_CACHE_MARKER, False):
+        setattr(cache, _OWNED_CACHE_MARKER, False)
+
+
 class KVCacheMemory(BaseActMemory):
     """
     Key-Value Cache Memory for activation memories.
@@ -58,6 +104,9 @@ def add(self, memories: list[KVCacheItem]) -> None:
             memories: List of KVCacheItem to add
         """
         for memory in memories:
+            # A stored cache belongs to the memory: it must not be treated as
+            # an owned copy, otherwise the model would append to it in place.
+            _mark_cache_as_stored(memory.memory)
             self.kv_cache_memories[memory.id] = memory
 
     def get_cache(self, cache_ids: list[str]) -> DynamicCache | None:
@@ -206,7 +255,9 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
 
         assert caches, "Need at least one cache"
         if len(caches) == 1:
-            return caches[0]
+            # Hand out an independent copy: the stored cache must stay read-only,
+            # otherwise the model appends the generated tokens to it in place.
+            return clone_dynamic_cache(caches[0])
 
         merged = DynamicCache()
 
@@ -252,9 +303,91 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
                 "DynamicCache object has neither 'layers' nor 'key_cache' attributes"
             )
 
+        # The merged cache is built from freshly concatenated tensors, so it
+        # shares nothing with the stored caches: it is an owned copy as well.
+        setattr(merged, _OWNED_CACHE_MARKER, True)
+
         return merged
 
 
+def _clone_tensor(tensor: Any) -> Any:
+    """Clone a cache tensor, keeping ``None`` as ``None``."""
+    if tensor is None:
+        return None
+    return tensor.clone()
+
+
+def _clone_cache_layer(layer: Any) -> Any:
+    """
+    Copy a single cache layer without sharing any tensor with the original.
+
+    The layer object itself is shallow-copied so that all non-tensor state
+    (e.g. the cumulative sequence length tracked by newer ``transformers``
+    versions) is preserved, while every key/value tensor is replaced by an
+    independent clone.
+    """
+    cloned = copy.copy(layer)
+    for attribute in ("keys", "key_cache", "values", "value_cache"):
+        tensor = getattr(layer, attribute, None)
+        if tensor is not None:
+            setattr(cloned, attribute, _clone_tensor(tensor))
+    return cloned
+
+
+def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
+    """
+    Return an independent copy of a ``DynamicCache``.
+
+    ``transformers`` appends the newly generated key/value tensors **in place**
+    to the cache object it is given (``DynamicLayer.update`` rebinds
+    ``keys``/``values`` on the very same layer object). When the cache handed
+    to the model is the object stored in the activation memory, every chat turn
+    permanently grows the stored activation memory, and the growth is written
+    back to disk on the next dump.
+
+    The returned cache shares no tensor with the original, so it can safely be
+    passed to the model. Both cache layouts are supported:
+
+    * ``transformers < 4.55``: ``key_cache`` / ``value_cache`` lists;
+    * ``transformers >= 4.55``: a ``layers`` list whose layers carry either
+      ``keys`` / ``values`` or ``key_cache`` / ``value_cache``.
+
+    The copy is marked as an owned copy (see :func:`is_owned_cache`), so the
+    model boundary does not copy it a second time when it is handed to a
+    model that appends to it in place.
+
+    Args:
+        cache: The cache to copy.
+
+    Returns:
+        DynamicCache: A copy that does not share any tensor with ``cache``.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        # The modern layout. The check is on the attribute itself and not on
+        # its content: an empty cache still has to be copied into an
+        # independent one, because the model appends its first layer to the
+        # cache it receives.
+        cloned = copy.copy(cache)
+        cloned.layers = [_clone_cache_layer(layer) for layer in layers]
+        setattr(cloned, _OWNED_CACHE_MARKER, True)
+        return cloned
+
+    key_cache = getattr(cache, "key_cache", None)
+    value_cache = getattr(cache, "value_cache", None)
+    if key_cache is not None or value_cache is not None:
+        # The legacy layout, where the tensors live in two flat lists.
+        cloned = copy.copy(cache)
+        if key_cache is not None:
+            cloned.key_cache = [_clone_tensor(tensor) for tensor in key_cache]
+        if value_cache is not None:
+            cloned.value_cache = [_clone_tensor(tensor) for tensor in value_cache]
+        setattr(cloned, _OWNED_CACHE_MARKER, True)
+        return cloned
+
+    raise AttributeError("DynamicCache object has neither 'layers' nor 'key_cache' attributes")
+
+
 def move_dynamic_cache_htod(dynamic_cache: DynamicCache, device: str) -> DynamicCache:
     """
     Move DynamicCache from CPU to GPU device.