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,36 @@ 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.
+
+        Args:
+            cache (DynamicCache): The cache provided by the caller.
+
+        Returns:
+            DynamicCache: A copy that does not share any tensor with ``cache``.
+        """
+        # 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
+
+        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
 
@@ -206,7 +208,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()
 
@@ -255,6 +259,78 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
         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``.
+
+    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]
+        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]
+        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.

tests/llms/test_hf.py

@@ -2,13 +2,46 @@
 
 from unittest.mock import MagicMock, patch
 
-import torch
+import pytest
 
 from transformers import DynamicCache
 
 from memos.configs.llm import HFLLMConfig, LLMConfigFactory
 from memos.llms.factory import LLMFactory
 from memos.llms.hf import HFLLM
+from tests.utils import dynamic_cache_seq_len, make_dynamic_cache
+
+
+# The KV cache is a torch structure; skip the whole module where torch is
+# not installed instead of failing the collection of the test suite.
+torch = pytest.importorskip("torch")
+
+
+class CacheGrowingModel:
+    """A stand-in for a transformers model that grows the cache in place.
+
+    Real transformers models append the newly generated key/value tensors to
+    the ``past_key_values`` object they are given (``DynamicLayer.update``
+    rebinds ``keys``/``values`` on the very same layer object). Handing the
+    stored activation cache to the model therefore grows the stored cache with
+    every turn.
+    """
+
+    def __init__(self, device: str = "cpu") -> None:
+        self.device = device
+        self.received_caches = []
+
+    def __call__(
+        self, input_ids=None, use_cache=False, past_key_values=None, return_dict=True, **kwargs
+    ):
+        self.received_caches.append(past_key_values)
+        if past_key_values is not None:
+            past_key_values.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+        output = MagicMock()
+        output.logits = torch.full((1, 1, 100), -1e9)
+        output.logits[0, 0, 7] = 1.0
+        output.past_key_values = past_key_values
+        return output
 
 
 @patch("transformers.AutoModelForCausalLM", MagicMock())
@@ -182,3 +215,52 @@ def test_kv_cache_generation_with_sampling(self):
         kv_cache = DynamicCache()
         resp = llm.generate([{"role": "user", "content": "Sampling"}], past_key_values=kv_cache)
         self.assertEqual(resp, self.standard_response)
+
+    def test_generate_does_not_mutate_caller_cache(self):
+        # Regression for #2301: the cache stored in the activation memory must
+        # not grow with the tokens generated for the current turn.
+        config = HFLLMConfig(
+            model_name_or_path="qwen3:0.6b",
+            max_tokens=5,
+            do_sample=False,
+            add_generation_prompt=True,
+        )
+        llm = self._create_llm(config)
+        model = CacheGrowingModel()
+        llm.model = model
+
+        cache = make_dynamic_cache(seq_len=3)
+        resp = llm.generate([{"role": "user", "content": "Hello"}], past_key_values=cache)
+
+        self.assertEqual(resp, self.standard_response)
+        # The model received a copy, not the caller's cache.
+        self.assertIsNot(model.received_caches[0], cache)
+        # The copy did grow, i.e. the model really appends in place.
+        self.assertGreater(dynamic_cache_seq_len(model.received_caches[0]), 3)
+        # The caller's cache still holds only the original tokens.
+        self.assertEqual(dynamic_cache_seq_len(cache), 3)
+
+    def test_generate_stream_does_not_mutate_caller_cache(self):
+        # Same regression for the streaming path.
+        config = HFLLMConfig(
+            model_name_or_path="qwen3:0.6b",
+            max_tokens=5,
+            do_sample=False,
+            add_generation_prompt=True,
+        )
+        llm = self._create_llm(config)
+        model = CacheGrowingModel()
+        llm.model = model
+
+        cache = make_dynamic_cache(seq_len=3)
+        chunks = list(
+            llm.generate_stream([{"role": "user", "content": "Hello"}], past_key_values=cache)
+        )
+
+        self.assertTrue(chunks)
+        # The model received a copy, not the caller's cache.
+        self.assertIsNot(model.received_caches[0], cache)
+        # The copy did grow, i.e. the model really appends in place.
+        self.assertGreater(dynamic_cache_seq_len(model.received_caches[0]), 3)
+        # The caller's cache still holds only the original tokens.
+        self.assertEqual(dynamic_cache_seq_len(cache), 3)