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.