From f5d6e15935de88408018bcb093d111bc31805912 Mon Sep 17 00:00:00 2001
From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:56:09 +0800
Subject: [PATCH 1/2] fix(activation): keep the stored KV cache read-only
 during generation (#2301)

The activation memory stored in `KVCacheMemory` was handed to the model as
the very same object, and `transformers` appends the newly generated
key/value tensors to that object **in place** (`DynamicLayer.update`
rebinds `keys`/`values` on the same layer object).

So the stored activation memory grew with the prompt and the generated
tokens of every turn, and `ActivationMemoryManager` wrote the growth back
to disk on the next dump. The stored cache was supposed to represent a
fixed set of memories, not the conversation.

Three hops were involved:

1. `MemOS.chat` (`mem_os/core.py`), `mem_os/main.py`, `mem_chat/simple.py`
   and the scheduler analyzer pass the stored `kv_cache.memory` straight
   into `HFLLM.generate(past_key_values=...)`;
2. `HFLLM._prefill` forwards it to `model(past_key_values=...)`, and
   transformers appends the new tokens' K/V in place;
3. `KVCacheMemory._concat_caches` returns `caches[0]` unchanged for a
   single id, so the `get_cache()` merge path hands out the stored object
   too.

This commit treats stored caches as read-only by construction, at the two
boundaries where a cache leaves the store or enters the model:

- `clone_dynamic_cache()` (new, `memories/activation/kv.py`): an
  independent copy with cloned K/V tensors. The modern layout
  (`layers`, transformers >= 4.55) is preferred whenever the attribute
  exists, so that the copy never shares its layer list with the
  original, even for an empty cache that the model fills with its first
  layer. The legacy layout (`key_cache`/`value_cache` lists, transformers
  < 4.55) is copied as well. On a layer object, `keys`/`values` and
  `key_cache`/`value_cache` are copied independently of each other, so a
  layer that populates only one side does not lose data. The layer
  objects are shallow-copied, so all non-tensor state (for example the
  initialization flag that `DynamicLayer.get_seq_length()` relies on) is
  preserved.
- `KVCacheMemory._concat_caches`: the single-cache case now returns a
  clone instead of the stored object.
- `HFLLM.generate` / `HFLLM.generate_stream`: clone the incoming
  `past_key_values` once at the boundary. This covers all four call
  sites without touching them; the cost is one cache copy per generate
  call, the same order as the prefill work itself.

No production dependency, configuration, schema or API change.

Tests:

- `tests/llms/test_hf.py`: `test_generate_does_not_mutate_caller_cache`
  and `test_generate_stream_does_not_mutate_caller_cache` use a stand-in
  model that appends K/V in place, and assert that the model receives a
  copy and that the caller's cache still holds only the original tokens.
  Both fail before the fix (the model receives the caller's cache
  itself) and pass after it.
- `tests/memories/activation/test_kv.py`:
  `test_clone_dynamic_cache_does_not_share_tensors`,
  `test_get_cache_single_item_returns_independent_copy`,
  `test_get_cache_merge_does_not_alias_inputs`,
  `test_clone_dynamic_cache_handles_empty_cache`,
  `test_clone_dynamic_cache_prefers_the_layers_layout`,
  `test_clone_dynamic_cache_handles_legacy_layout` and
  `test_clone_dynamic_cache_copies_layers_with_legacy_attribute_names`
  cover the store boundary and both cache layouts.
- `tests/utils.py`: shared helpers to build a cache and to read it
  independently of the cache layout.

The layout-independent helpers also replace the `key_cache`-only helper in
`tests/memories/activation/test_kv.py`, which cannot build a populated
cache at all with transformers >= 4.55.

The two touched test modules now import torch through
`pytest.importorskip("torch")`, so that a test environment without torch
skips them instead of failing during collection.

Fixes #2301.
---
 src/memos/llms/hf.py                 |  33 ++++-
 src/memos/memories/activation/kv.py  |  78 +++++++++++-
 tests/llms/test_hf.py                |  84 ++++++++++++-
 tests/memories/activation/test_kv.py | 174 +++++++++++++++++++++++++--
 tests/utils.py                       |  82 +++++++++++++
 5 files changed, 438 insertions(+), 13 deletions(-)

diff --git a/src/memos/llms/hf.py b/src/memos/llms/hf.py
index 0dd841c1a..ad0363187 100644
--- a/src/memos/llms/hf.py
+++ b/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:
         """
diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py
index 1981b958f..166f61c90 100644
--- a/src/memos/memories/activation/kv.py
+++ b/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.
diff --git a/tests/llms/test_hf.py b/tests/llms/test_hf.py
index 375bf2247..d5265e2f5 100644
--- a/tests/llms/test_hf.py
+++ b/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)
diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py
index 6490d687f..2a3b5f645 100644
--- a/tests/memories/activation/test_kv.py
+++ b/tests/memories/activation/test_kv.py
@@ -1,13 +1,23 @@
 from unittest.mock import MagicMock
 
 import pytest
-import torch
 
 from transformers import DynamicCache
 
 from memos.configs.memory import KVCacheMemoryConfig
 from memos.memories.activation.item import KVCacheItem
-from memos.memories.activation.kv import KVCacheMemory
+from memos.memories.activation.kv import KVCacheMemory, clone_dynamic_cache
+from tests.utils import (
+    dynamic_cache_key_tensors,
+    dynamic_cache_layer_count,
+    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")
 
 
 @pytest.fixture
@@ -35,10 +45,7 @@ def kv_memory(dummy_config):
 
 def make_filled_cache():
     # Create a DynamicCache with at least one dummy tensor layer
-    cache = DynamicCache()
-    cache.key_cache.append(torch.zeros(1, 2, 3))
-    cache.value_cache.append(torch.zeros(1, 2, 3))
-    return cache
+    return make_dynamic_cache(seq_len=3)
 
 
 def test_extract_and_add_and_get(kv_memory):
@@ -58,9 +65,8 @@ def test_get_cache_merge(kv_memory):
     kv_memory.add([item1, item2])
     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
+    # Check the number of layers in merged cache
+    assert dynamic_cache_layer_count(merged) == 1
 
 
 def test_delete_and_get_all(kv_memory):
@@ -84,3 +90,153 @@ class DummyTextualMemory:
     item = kv_memory.from_textual_memory(DummyTextualMemory())
     assert isinstance(item, KVCacheItem)
     assert item.metadata["bar"] == 1
+
+
+def test_clone_dynamic_cache_does_not_share_tensors():
+    # A clone must not share any tensor with the original: transformers appends
+    # the generated tokens to the cache object in place.
+    cache = make_filled_cache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert dynamic_cache_layer_count(cloned) == dynamic_cache_layer_count(cache)
+
+    original_keys = dynamic_cache_key_tensors(cache)[0]
+    cloned_keys = dynamic_cache_key_tensors(cloned)[0]
+    assert cloned_keys is not original_keys
+    assert cloned_keys.data_ptr() != original_keys.data_ptr()
+    assert torch.equal(cloned_keys, original_keys)
+
+    # The copy still reports the same number of cached tokens, i.e. the layer
+    # state (such as the initialization flag) survives the copy.
+    assert cloned.get_seq_length() == cache.get_seq_length()
+
+    # Appending to the clone (what a model does) leaves the original untouched.
+    cloned.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+    assert dynamic_cache_seq_len(cloned) == 4
+    assert dynamic_cache_seq_len(cache) == 3
+    assert cloned.get_seq_length() == 4
+    assert cache.get_seq_length() == 3
+
+
+def test_get_cache_single_item_returns_independent_copy(kv_memory):
+    # The stored cache must never be handed out: the model appends the
+    # generated tokens to the cache it receives, in place.
+    item = KVCacheItem(memory=make_filled_cache())
+    kv_memory.add([item])
+
+    cached = kv_memory.get_cache([item.id])
+
+    assert cached is not None
+    assert cached is not item.memory
+    assert dynamic_cache_seq_len(cached) == 3
+
+    cached.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_seq_len(cached) == 4
+    assert dynamic_cache_seq_len(item.memory) == 3
+
+
+def test_clone_dynamic_cache_handles_empty_cache():
+    # An empty cache has to be copied into an independent empty cache: the
+    # model appends its first layer to the cache it receives, and that append
+    # must not reach the stored cache.
+    cache = DynamicCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert dynamic_cache_layer_count(cloned) == 0
+    assert dynamic_cache_layer_count(cache) == 0
+
+    cloned.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_layer_count(cloned) == 1
+    assert dynamic_cache_layer_count(cache) == 0
+    assert cache.get_seq_length() == 0
+
+
+def test_get_cache_merge_does_not_alias_inputs(kv_memory):
+    # Merging concatenates the inputs into new tensors; the stored caches
+    # stay untouched.
+    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 merged is not None
+    assert dynamic_cache_layer_count(merged) == 1
+    # Two caches of 3 tokens each are merged into a single 6-token layer.
+    assert dynamic_cache_seq_len(merged) == 6
+
+    merged.update(torch.zeros(1, 2, 1, 4), torch.ones(1, 2, 1, 4), 0)
+
+    assert dynamic_cache_seq_len(item1.memory) == 3
+    assert dynamic_cache_seq_len(item2.memory) == 3
+
+
+def test_clone_dynamic_cache_handles_legacy_layout():
+    # transformers < 4.55 keeps the tensors in the flat "key_cache" and
+    # "value_cache" lists instead of per-layer objects.
+    class LegacyCache:
+        def __init__(self):
+            self.key_cache = [torch.zeros(1, 2, 3, 4)]
+            self.value_cache = [torch.ones(1, 2, 3, 4)]
+
+    cache = LegacyCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.key_cache[0] is not cache.key_cache[0]
+    assert cloned.key_cache[0].data_ptr() != cache.key_cache[0].data_ptr()
+    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])
+
+
+def test_clone_dynamic_cache_prefers_the_layers_layout():
+    # A cache that exposes the modern "layers" layout next to the legacy
+    # "key_cache"/"value_cache" aliases must be copied through the modern
+    # layout: the model appends its layers to "layers", so that list must not
+    # be shared with the original.
+    class CacheWithLegacyAliases:
+        def __init__(self):
+            self.layers = []
+            self.key_cache = []
+            self.value_cache = []
+
+        def update(self, key_states, value_states, layer_idx):
+            self.layers.append(layer_idx)
+
+    cache = CacheWithLegacyAliases()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.layers is not cache.layers
+
+    cloned.update(None, None, 0)
+
+    assert cache.layers == []
+
+
+def test_clone_dynamic_cache_copies_layers_with_legacy_attribute_names():
+    # Some transformers versions expose "key_cache"/"value_cache" on the layer
+    # objects instead of "keys"/"values"; both variants have to be copied.
+    class LegacyLayer:
+        def __init__(self):
+            self.key_cache = torch.zeros(1, 2, 3, 4)
+            self.value_cache = torch.ones(1, 2, 3, 4)
+
+    class LayeredCache:
+        def __init__(self):
+            self.layers = [LegacyLayer()]
+
+    cache = LayeredCache()
+    cloned = clone_dynamic_cache(cache)
+
+    assert cloned is not cache
+    assert cloned.layers[0] is not cache.layers[0]
+    assert cloned.layers[0].key_cache.data_ptr() != cache.layers[0].key_cache.data_ptr()
+    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)
diff --git a/tests/utils.py b/tests/utils.py
index ec8a32799..0f1cf4842 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -9,6 +9,88 @@
 from pydantic.aliases import PydanticUndefined
 
 
+def make_dynamic_cache(seq_len: int = 3, num_layers: int = 1):
+    """
+    Build a ``DynamicCache`` with populated layers.
+
+    ``transformers`` stores the key/value tensors either in the legacy
+    ``key_cache``/``value_cache`` lists (< 4.55) or in per-layer objects
+    (``layers``, >= 4.55). Building the cache through the public ``update()``
+    API works for both layouts.
+
+    Args:
+        seq_len: Number of cached tokens per layer.
+        num_layers: Number of layers to populate.
+
+    Returns:
+        DynamicCache: A cache with ``num_layers`` populated layers.
+    """
+    import torch
+
+    from transformers import DynamicCache
+
+    cache = DynamicCache()
+    for layer_idx in range(num_layers):
+        cache.update(
+            torch.zeros(1, 2, seq_len, 4),
+            torch.ones(1, 2, seq_len, 4),
+            layer_idx,
+        )
+    return cache
+
+
+def dynamic_cache_key_tensors(cache: Any) -> list[Any]:
+    """
+    Return the key tensor of every layer of a cache, for either cache layout.
+
+    Args:
+        cache: The cache to inspect.
+
+    Returns:
+        list: One key tensor per layer.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        tensors = []
+        for layer in layers:
+            tensor = getattr(layer, "keys", None)
+            if tensor is None:
+                tensor = getattr(layer, "key_cache", None)
+            tensors.append(tensor)
+        return tensors
+    return list(cache.key_cache)
+
+
+def dynamic_cache_layer_count(cache: Any) -> int:
+    """
+    Return the number of layers of a cache, for either cache layout.
+
+    Args:
+        cache: The cache to inspect.
+
+    Returns:
+        int: Number of layers.
+    """
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        return len(layers)
+    return len(cache.key_cache)
+
+
+def dynamic_cache_seq_len(cache: Any, layer_idx: int = 0) -> int:
+    """
+    Return the number of cached tokens held by one layer of a cache.
+
+    Args:
+        cache: The cache to inspect.
+        layer_idx: Index of the layer to look at.
+
+    Returns:
+        int: Number of cached tokens.
+    """
+    return dynamic_cache_key_tensors(cache)[layer_idx].shape[-2]
+
+
 def check_module_base_class(cls: Any) -> None:
     """
     General function to test the correctness of an abstract base class.

From 55d2d84269dcb5ae3c598f1861adfa2ce40e32fa Mon Sep 17 00:00:00 2001
From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com>
Date: Tue, 3 Nov 2026 11:29:51 +0800
Subject: [PATCH 2/2] 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.
---
 src/memos/llms/hf.py                 |  15 +++-
 src/memos/memories/activation/kv.py  |  57 +++++++++++++++
 tests/llms/test_hf.py                | 104 ++++++++++++++++++++++++++-
 tests/memories/activation/test_kv.py |  60 +++++++++++++++-
 4 files changed, 232 insertions(+), 4 deletions(-)

diff --git a/src/memos/llms/hf.py b/src/memos/llms/hf.py
index ad0363187..fb9d1d419 100644
--- a/src/memos/llms/hf.py
+++ b/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)
 
diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py
index 166f61c90..1d0879fc5 100644
--- a/src/memos/memories/activation/kv.py
+++ b/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")
diff --git a/tests/llms/test_hf.py b/tests/llms/test_hf.py
index d5265e2f5..7363bfc16 100644
--- a/tests/llms/test_hf.py
+++ b/tests/llms/test_hf.py
@@ -9,7 +9,12 @@
 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
@@ -17,6 +22,30 @@
 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)
diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py
index 2a3b5f645..87e10219d 100644
--- a/tests/memories/activation/test_kv.py
+++ b/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)