fix(activation): copy the KV cache only once per generation

The activation memory hands out an owned copy of the stored cache
(`get_cache()`), and `HFLLM.generate` copied that copy again before it
was handed to the model. A caller that follows the public API therefore
paid for two full copies of every K/V tensor and kept three sets alive
during the generation: the stored cache, the copy of the copy, and the
one the model writes into.

Ownership is now explicit. A cache that leaves the activation memory is
an *owned copy*: it shares no tensor with the stored cache, so it may be
handed to a model that appends to it in place. `clone_dynamic_cache()`
and the merge path of `_concat_caches()` mark the cache they return, and
the model boundary skips the copy for a cache that is already owned. A
cache that is stored back into the memory loses the marker, so the model
still copies it before it appends to it.

The stored cache stays read-only in every case: the model never writes
into a cache that aliases it, and the number of K/V tensor sets alive
during a generation drops from three to two.

Tests:

- `tests/llms/test_hf.py::test_generation_from_get_cache_copies_the_kv_tensors_only_once`
  walks the public path (`get_cache()` -> `generate()`), counts the copies
  (exactly one) and pins that the model received the cache the memory
  handed out instead of a second copy of it.
- `tests/llms/test_hf.py::test_generation_from_get_cache_keeps_the_stored_cache_read_only`
  pins that the stored cache is not rewritten, also when an owned copy is
  stored back into the memory.
- `tests/memories/activation/test_kv.py` pins the ownership marker on the
  single-cache and on the merge path, and that a stored cache is never an
  owned copy.

Fixes #2301.

Author: RerankerGuo

src/memos/llms/hf.py

@@ -121,15 +121,26 @@ def _clone_cache(cache: DynamicCache) -> DynamicCache:
         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 copy that does not share any tensor with ``cache``.
+            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
+        from memos.memories.activation.kv import clone_dynamic_cache, is_owned_cache
+
+        if is_owned_cache(cache):
+            return cache
 
         return clone_dynamic_cache(cache)
 

src/memos/memories/activation/kv.py

@@ -15,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.
@@ -60,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:
@@ -256,6 +303,10 @@ 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
 
 
@@ -301,6 +352,10 @@ def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
     * ``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.
 
@@ -315,6 +370,7 @@ def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
         # 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)
@@ -326,6 +382,7 @@ def clone_dynamic_cache(cache: DynamicCache) -> DynamicCache:
             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")

tests/llms/test_hf.py

@@ -9,14 +9,43 @@
 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
+from memos.memories.activation.kv import clone_dynamic_cache
+from tests.utils import (
+    dynamic_cache_key_tensors,
+    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")
 
 
+def make_activation_memory(cache):
+    """Build a ``KVCacheMemory`` that stores a single KV cache.
+
+    Args:
+        cache: The cache to store in the memory.
+
+    Returns:
+        tuple: The memory and the stored ``KVCacheItem``.
+    """
+    from memos.configs.memory import KVCacheMemoryConfig
+    from memos.memories.activation.item import KVCacheItem
+    from memos.memories.activation.kv import KVCacheMemory
+
+    config = MagicMock(spec=KVCacheMemoryConfig)
+    config.extractor_llm = MagicMock()
+    config.memory_filename = "test_kv_cache.pkl"
+    with patch("memos.llms.factory.LLMFactory.from_config", return_value=MagicMock()):
+        memory = KVCacheMemory(config)
+
+    item = KVCacheItem(memory=cache)
+    memory.add([item])
+    return memory, item
+
+
 class CacheGrowingModel:
     """A stand-in for a transformers model that grows the cache in place.
 
@@ -264,3 +293,76 @@ def test_generate_stream_does_not_mutate_caller_cache(self):
         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_generation_from_get_cache_copies_the_kv_tensors_only_once(self):
+        # The public API path: the activation memory hands out a cache with
+        # get_cache(), the caller passes it to generate(). The stored cache
+        # must stay read-only, and the K/V tensors must not be copied twice
+        # along the way -- one copy would already keep three full sets of
+        # tensors alive (stored + copy 1 + copy 2).
+        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
+
+        stored_cache = make_dynamic_cache(seq_len=3)
+        memory, item = make_activation_memory(stored_cache)
+
+        with patch(
+            "memos.memories.activation.kv.clone_dynamic_cache", wraps=clone_dynamic_cache
+        ) as clone_spy:
+            cached = memory.get_cache([item.id])
+            resp = llm.generate([{"role": "user", "content": "Hello"}], past_key_values=cached)
+
+        self.assertEqual(resp, self.standard_response)
+        # The whole path copies the K/V tensors exactly once.
+        self.assertEqual(clone_spy.call_count, 1)
+        # The model received the cache the memory handed out, not a second copy.
+        self.assertIs(model.received_caches[0], cached)
+        # The copy the model wrote into does not alias the stored tensors.
+        self.assertNotEqual(
+            dynamic_cache_key_tensors(model.received_caches[0])[0].data_ptr(),
+            dynamic_cache_key_tensors(item.memory)[0].data_ptr(),
+        )
+        # The copy did grow, i.e. the model really appends in place.
+        self.assertGreater(dynamic_cache_seq_len(model.received_caches[0]), 3)
+        # The stored cache was not rewritten: it still holds only the
+        # original tokens.
+        self.assertEqual(dynamic_cache_seq_len(item.memory), 3)
+        self.assertEqual(dynamic_cache_seq_len(stored_cache), 3)
+
+    def test_generation_from_get_cache_keeps_the_stored_cache_read_only(self):
+        # The stored cache stays read-only even when it is handed out a second
+        # time and when a caller stores an already owned copy back.
+        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
+
+        stored_cache = make_dynamic_cache(seq_len=3)
+        memory, item = make_activation_memory(stored_cache)
+
+        cached = memory.get_cache([item.id])
+        llm.generate([{"role": "user", "content": "Hello"}], past_key_values=cached)
+
+        # The stored cache kept its original length.
+        self.assertEqual(dynamic_cache_seq_len(item.memory), 3)
+        self.assertEqual(dynamic_cache_seq_len(stored_cache), 3)
+
+        # A cache that is stored back into the memory is owned by the memory
+        # again: the model must copy it before it appends to it.
+        memory.add([item])
+        llm.generate([{"role": "user", "content": "Hello"}], past_key_values=item.memory)
+
+        self.assertIsNot(model.received_caches[-1], item.memory)
+        self.assertEqual(dynamic_cache_seq_len(item.memory), 3)

tests/memories/activation/test_kv.py

@@ -6,7 +6,11 @@
 
 from memos.configs.memory import KVCacheMemoryConfig
 from memos.memories.activation.item import KVCacheItem
-from memos.memories.activation.kv import KVCacheMemory, clone_dynamic_cache
+from memos.memories.activation.kv import (
+    KVCacheMemory,
+    clone_dynamic_cache,
+    is_owned_cache,
+)
 from tests.utils import (
     dynamic_cache_key_tensors,
     dynamic_cache_layer_count,
@@ -192,6 +196,9 @@ def __init__(self):
     assert torch.equal(cloned.key_cache[0], cache.key_cache[0])
     assert cloned.value_cache[0].data_ptr() != cache.value_cache[0].data_ptr()
     assert torch.equal(cloned.value_cache[0], cache.value_cache[0])
+    # The copy of a legacy cache is an owned copy as well.
+    assert is_owned_cache(cloned)
+    assert not is_owned_cache(cache)
 
 
 def test_clone_dynamic_cache_prefers_the_layers_layout():
@@ -240,3 +247,54 @@ def __init__(self):
     assert torch.equal(cloned.layers[0].key_cache, cache.layers[0].key_cache)
     assert cloned.layers[0].value_cache.data_ptr() != cache.layers[0].value_cache.data_ptr()
     assert torch.equal(cloned.layers[0].value_cache, cache.layers[0].value_cache)
+
+
+def test_clone_dynamic_cache_hands_out_an_owned_copy():
+    # A clone is an owned copy: it shares no tensor with the original, so the
+    # caller may hand it to a model that appends to it in place. The model
+    # boundary uses that ownership to avoid copying it a second time.
+    cache = make_filled_cache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert is_owned_cache(cloned)
+    assert not is_owned_cache(cache)
+
+
+def test_get_cache_hands_out_an_owned_copy(kv_memory):
+    # The cache the memory hands out is an owned copy; the stored cache is not.
+    item = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item])
+
+    cached = kv_memory.get_cache([item.id])
+
+    assert is_owned_cache(cached)
+    assert not is_owned_cache(item.memory)
+
+
+def test_get_cache_merge_hands_out_an_owned_copy(kv_memory):
+    # The merged cache is built from freshly concatenated tensors, so it is an
+    # owned copy as well.
+    item1 = KVCacheItem(memory=make_filled_cache())
+    item2 = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item1, item2])
+
+    merged = kv_memory.get_cache([item1.id, item2.id])
+
+    assert is_owned_cache(merged)
+    assert not is_owned_cache(item1.memory)
+    assert not is_owned_cache(item2.memory)
+
+
+def test_stored_cache_is_never_an_owned_copy(kv_memory):
+    # A cache that is stored in the memory belongs to the memory: storing an
+    # already owned copy back must drop the ownership, otherwise the model
+    # would append to the stored cache in place.
+    item = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item])
+
+    cached = kv_memory.get_cache([item.id])
+    assert is_owned_cache(cached)
+
+    kv_memory.add([KVCacheItem(memory=cached, metadata=item.metadata)])
+
+    assert not is_owned_cache(cached)