fix: stop activation KV cache from growing every turn

src/memos/llms/hf.py

@@ -82,7 +82,12 @@ 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)
+            from memos.memories.activation.kv import clone_dynamic_cache
+
+            # The model appends new K/V tensors to the cache it receives, so
+            # hand it a clone and keep the caller's cache (e.g. a stored
+            # activation memory) unchanged by this call.
+            return self._generate_with_cache(prompt, clone_dynamic_cache(past_key_values), **kwargs)
 
     def generate_stream(
         self, messages: MessageList, past_key_values: DynamicCache | None = None, **kwargs
@@ -102,7 +107,11 @@ 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)
+            from memos.memories.activation.kv import clone_dynamic_cache
+
+            yield from self._generate_with_cache_stream(
+                prompt, clone_dynamic_cache(past_key_values)
+            )
 
     def _generate_full(self, prompt: str, **kwargs) -> str:
         """

src/memos/memories/activation/kv.py

@@ -1,3 +1,4 @@
+import copy
 import os
 import pickle
 
@@ -206,7 +207,10 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
 
         assert caches, "Need at least one cache"
         if len(caches) == 1:
-            return caches[0]
+            # Return a copy: the stored cache must never be handed out by
+            # reference, because generation appends new K/V tensors to the
+            # cache object it receives and would grow the store every turn.
+            return clone_dynamic_cache(caches[0])
 
         merged = DynamicCache()
 
@@ -248,13 +252,89 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
                 merged.value_cache.append(torch.cat(vals, dim=-2))
 
         else:
-            raise AttributeError(
-                "DynamicCache object has neither 'layers' nor 'key_cache' attributes"
-            )
+            raise TypeError("DynamicCache object has neither 'layers' nor 'key_cache' attributes")
 
         return merged
 
 
+def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
+    """
+    Return an independent copy of a DynamicCache with cloned K/V tensors.
+
+    Generation mutates the cache object it receives in place, so a stored cache
+    must never be handed to a model by reference — hand out a clone instead.
+    Compatible with both old (key_cache/value_cache) and new (layers) structures.
+    """
+    import torch
+
+    cloned = DynamicCache()
+
+    if hasattr(cache, "layers"):
+        cloned.layers = []
+        for layer in cache.layers:
+            # Avoid invoking a layer constructor: modern transformers layers
+            # such as DynamicSlidingWindowLayer require constructor metadata.
+            new_layer = copy.copy(layer)
+            layer_attrs = vars(layer)
+            # Preserve layer state and clone every tensor, including K/V tensors.
+            kv_attrs = {"keys", "values", "key_cache", "value_cache"}
+            for attr, value in layer_attrs.items():
+                if attr in kv_attrs:
+                    continue
+                setattr(
+                    new_layer,
+                    attr,
+                    value.clone() if isinstance(value, torch.Tensor) else copy.deepcopy(value),
+                )
+            # transformers>=4.56 layers expose keys/values, but some versions
+            # instead carry per-layer key_cache/value_cache (see
+            # move_dynamic_cache_htod); a clone that skips one shape would
+            # silently return a content-empty layer.
+            # Select one naming scheme, matching move_dynamic_cache_htod's
+            # precedence, while retaining independent guards for asymmetric
+            # test doubles and cache layers.
+            has_per_layer_cache = any(
+                getattr(layer, name, None) is not None for name in ("key_cache", "value_cache")
+            )
+            if has_per_layer_cache:
+                if "keys" in layer_attrs:
+                    new_layer.keys = None
+                if "values" in layer_attrs:
+                    new_layer.values = None
+                if getattr(layer, "key_cache", None) is not None:
+                    new_layer.key_cache = layer.key_cache.clone()
+                if getattr(layer, "value_cache", None) is not None:
+                    new_layer.value_cache = layer.value_cache.clone()
+            else:
+                if "key_cache" in layer_attrs:
+                    new_layer.key_cache = None
+                if "value_cache" in layer_attrs:
+                    new_layer.value_cache = None
+                if getattr(layer, "keys", None) is not None:
+                    new_layer.keys = layer.keys.clone()
+                if getattr(layer, "values", None) is not None:
+                    new_layer.values = layer.values.clone()
+            cloned.layers.append(new_layer)
+    elif hasattr(cache, "key_cache"):
+        # Legacy DynamicCache keeps generation state such as _seen_tokens on
+        # the cache itself.  Keep that state independent of the stored cache;
+        # key/value lists are populated from cloned tensors below.
+        for attr, value in vars(cache).items():
+            if attr not in {"key_cache", "value_cache"}:
+                setattr(
+                    cloned,
+                    attr,
+                    value.clone() if isinstance(value, torch.Tensor) else copy.deepcopy(value),
+                )
+        for keys, values in zip(cache.key_cache, cache.value_cache, strict=True):
+            cloned.key_cache.append(keys.clone() if keys is not None else None)
+            cloned.value_cache.append(values.clone() if values is not None else None)
+    else:
+        raise TypeError("DynamicCache object has neither 'layers' nor 'key_cache' attributes")
+
+    return cloned
+
+
 def move_dynamic_cache_htod(dynamic_cache: DynamicCache, device: str) -> DynamicCache:
     """
     Move DynamicCache from CPU to GPU device.

tests/cache_helpers.py

@@ -0,0 +1,63 @@
+import pytest
+import torch
+
+from transformers import DynamicCache
+
+
+def make_filled_cache():
+    cache = DynamicCache()
+    keys = torch.zeros(1, 2, 3, 4) if hasattr(cache, "layers") else torch.zeros(1, 2, 3)
+    values = torch.zeros_like(keys)
+    cache.update(keys, values, layer_idx=0)
+    return cache
+
+
+def cache_keys(cache, layer_idx=0):
+    if hasattr(cache, "layers"):
+        return cache.layers[layer_idx].keys
+    return cache.key_cache[layer_idx]
+
+
+def cache_values(cache, layer_idx=0):
+    if hasattr(cache, "layers"):
+        return cache.layers[layer_idx].values
+    return cache.value_cache[layer_idx]
+
+
+def set_cache_keys(cache, value, layer_idx=0):
+    if hasattr(cache, "layers"):
+        cache.layers[layer_idx].keys = value
+    else:
+        cache.key_cache[layer_idx] = value
+
+
+def cache_layer_count(cache):
+    if hasattr(cache, "layers"):
+        return len(cache.layers)
+    return len(cache.key_cache)
+
+
+def make_real_hybrid_cache(populate=True):
+    if not hasattr(DynamicCache(), "layers"):
+        pytest.skip("requires transformers >=4.56")
+
+    class HybridConfig:
+        num_hidden_layers = 2
+        sliding_window = 4
+
+        def __init__(self):
+            self.layer_types = ["full_attention", "sliding_attention"]
+
+        def get_text_config(self):
+            return self
+
+    try:
+        cache = DynamicCache(config=HybridConfig())
+        if populate:
+            keys = torch.zeros(1, 2, 3, 4)
+            values = torch.zeros(1, 2, 3, 4)
+            cache.update(keys, values, layer_idx=0)
+            cache.update(keys, values, layer_idx=1)
+    except TypeError:
+        pytest.skip("DynamicCache(config=...) is not supported")
+    return cache

tests/llms/test_hf.py

@@ -9,6 +9,9 @@
 from memos.configs.llm import HFLLMConfig, LLMConfigFactory
 from memos.llms.factory import LLMFactory
 from memos.llms.hf import HFLLM
+from tests.cache_helpers import cache_keys as _cache_keys
+from tests.cache_helpers import cache_values as _cache_values
+from tests.cache_helpers import make_filled_cache as _make_filled_cache
 
 
 @patch("transformers.AutoModelForCausalLM", MagicMock())
@@ -182,3 +185,56 @@ 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_with_cache_does_not_mutate_caller_cache(self):
+        """Regression for issue #2301: generation must not append K/V tensors
+        into the caller's stored cache (activation memory grew every turn)."""
+        config = HFLLMConfig(
+            model_name_or_path="qwen3:0.6b",
+            temperature=0.7,
+            max_tokens=3,
+            do_sample=True,
+            add_generation_prompt=True,
+        )
+        llm = self._create_llm(config)
+
+        kv_cache = _make_filled_cache()
+        original_key_shape = _cache_keys(kv_cache).shape
+        original_value_shape = _cache_values(kv_cache).shape
+        captured = {}
+
+        def forward(*args, **kwargs):
+            # transformers appends the new tokens' K/V to the cache in place.
+            # _prefill always passes the cache by keyword; .get keeps the mock
+            # resilient to an explicit-None caller without inventing a
+            # positional call shape.
+            kv = kwargs.get("past_key_values")
+            self.assertIsNotNone(kv, "forward() called without past_key_values")
+            captured["kv"] = kv
+            if hasattr(kv, "layers"):
+                kv.layers[0].keys = torch.cat([kv.layers[0].keys, torch.ones(1, 2, 1, 4)], dim=-2)
+                kv.layers[0].values = torch.cat(
+                    [kv.layers[0].values, torch.ones(1, 2, 1, 4)], dim=-2
+                )
+            else:
+                kv.key_cache[0] = torch.cat([kv.key_cache[0], torch.ones(1, 1, 3)], dim=-2)
+                kv.value_cache[0] = torch.cat([kv.value_cache[0], torch.ones(1, 1, 3)], dim=-2)
+            out = MagicMock()
+            # Deterministic non-EOS argmax so the loop runs all max_tokens turns
+            # instead of sometimes sampling eos_token_id (2) on the first step.
+            logits = torch.full((1, 1, 100), -1e9)
+            logits[0, 0, 10] = 0.0
+            out.logits = logits
+            out.past_key_values = kv
+            return out
+
+        self.mock_model.side_effect = forward
+        try:
+            llm.generate([{"role": "user", "content": "Hi"}], past_key_values=kv_cache)
+        finally:
+            self.mock_model.side_effect = None
+
+        self.assertEqual(_cache_keys(kv_cache).shape, original_key_shape)
+        self.assertEqual(_cache_values(kv_cache).shape, original_value_shape)
+        self.assertIsNotNone(captured.get("kv"))
+        self.assertIsNot(captured["kv"], kv_cache)