From 893d4b925aebbe1f7fcb1b05d53925498262d9da Mon Sep 17 00:00:00 2001 From: jeojdi1 Date: Thu, 27 Aug 2026 18:30:13 -0400 Subject: [PATCH] fix: stop merged KV caches being silently discarded on transformers >= 4.57 `_concat_caches` builds each merged layer with `layer_cls()` and then assigns `.keys` / `.values` directly. A layer constructed that way is never marked initialized, and on transformers >= 4.57 `DynamicLayer.get_seq_length()` short-circuits on that flag: if not self.is_initialized or self.keys.numel() == 0: return 0 So a correctly concatenated cache reports a length of 0, the model treats it as empty, and every merged token is dropped on the first forward pass. No error is raised and no warning is emitted -- the activation memory simply has no effect. Initialize the layer through its public `lazy_initialization` path before assigning, guarded by `hasattr` so the older `key_cache` layout is untouched. Also repairs the test fixture, which is what hid this. `make_filled_cache` appended to `cache.key_cache`, removed in 4.57, so `test_get_cache_merge` and `test_delete_and_get_all` failed with AttributeError on any current install and were being written off as version noise (see PR #2204's description). Rebuilding the fixture on the public `update` API makes those two tests pass again *and* makes them exercise the `layers` path, which the old fixture never did. Adds `test_concat_caches_preserves_seq_length`, which fails on the current code and passes with this change. --- src/memos/memories/activation/kv.py | 23 ++++++++++- tests/memories/activation/test_kv.py | 60 ++++++++++++++++++++++++---- 2 files changed, 74 insertions(+), 9 deletions(-) diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..9ee3c9f0f 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -232,8 +232,27 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache: keys = [c.layers[layer].keys for c in caches] vals = [c.layers[layer].values for c in caches] # single concat per layer - merged.layers[layer].keys = torch.cat(keys, dim=-2) - merged.layers[layer].values = torch.cat(vals, dim=-2) + merged_keys = torch.cat(keys, dim=-2) + merged_values = torch.cat(vals, dim=-2) + + # A layer constructed directly is not "initialized", and on + # transformers >= 4.57 ``DynamicLayer.get_seq_length()`` returns 0 + # for an uninitialized layer *regardless of the tensors it holds*: + # + # if not self.is_initialized or self.keys.numel() == 0: + # return 0 + # + # Assigning ``.keys``/``.values`` therefore leaves the merged cache + # reporting length 0, and the model treats it as empty and discards + # every merged token on the first forward pass -- silently, with no + # error. Initialize through the public path so the length is real. + merged_layer = merged.layers[layer] + if not getattr(merged_layer, "is_initialized", True) and hasattr( + merged_layer, "lazy_initialization" + ): + merged_layer.lazy_initialization(merged_keys, merged_values) + merged_layer.keys = merged_keys + merged_layer.values = merged_values # Check for old structure (key_cache) elif hasattr(caches[0], "key_cache"): diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 6490d687f..4437d1a64 100644 --- a/tests/memories/activation/test_kv.py +++ b/tests/memories/activation/test_kv.py @@ -33,11 +33,22 @@ def kv_memory(dummy_config): yield KVCacheMemory(dummy_config) -def make_filled_cache(): - # Create a DynamicCache with at least one dummy tensor layer +def make_filled_cache(seq_len: int = 3, n_layers: int = 1): + """Create a DynamicCache with dummy tensors, on any transformers version. + + This previously appended to ``cache.key_cache`` directly. That attribute was + removed in transformers 4.57 in favour of ``cache.layers``, so + ``test_get_cache_merge`` and ``test_delete_and_get_all`` failed with + ``AttributeError: 'DynamicCache' object has no attribute 'key_cache'`` on any + current install. Going through the public ``update`` API works on both the + old and new layouts -- and, unlike direct assignment, produces a properly + initialized cache, which is what makes the merge assertions meaningful. + """ cache = DynamicCache() - cache.key_cache.append(torch.zeros(1, 2, 3)) - cache.value_cache.append(torch.zeros(1, 2, 3)) + for layer_idx in range(n_layers): + k = torch.zeros(1, 2, seq_len, 4) + v = torch.zeros(1, 2, seq_len, 4) + cache.update(k, v, layer_idx) return cache @@ -58,9 +69,15 @@ 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 the merged cache. ``key_cache`` was removed + # in transformers 4.57; ``layers`` is the current layout, so accept either. + if hasattr(merged, "layers"): + assert len(merged.layers) == 1 + assert merged.layers[0].keys.shape[-2] == 6 + assert merged.layers[0].values.shape[-2] == 6 + else: + assert len(merged.key_cache) == 1 + assert len(merged.value_cache) == 1 def test_delete_and_get_all(kv_memory): @@ -84,3 +101,32 @@ class DummyTextualMemory: item = kv_memory.from_textual_memory(DummyTextualMemory()) assert isinstance(item, KVCacheItem) assert item.metadata["bar"] == 1 + + +def test_concat_caches_preserves_seq_length(kv_memory): + """A merged cache must report the summed length, not zero. + + Regression test for the merged cache being silently discarded. A layer built + by ``layer_cls()`` and populated by direct attribute assignment is never + marked initialized, and on transformers >= 4.57 + ``DynamicLayer.get_seq_length()`` short-circuits to 0 for an uninitialized + layer regardless of the tensors it holds. The merged cache therefore looked + empty and every merged token was dropped on the first forward pass, with no + error raised. + + Before the fix this asserts 0 == 7 and fails on transformers >= 4.57; it + passes on 4.56 and below, which is why it presented as a version nuisance. + """ + a = make_filled_cache(seq_len=3, n_layers=2) + b = make_filled_cache(seq_len=4, n_layers=2) + + merged = kv_memory._concat_caches([a, b]) + + assert merged.get_seq_length() == 7, ( + f"merged cache reports {merged.get_seq_length()} tokens, expected 7 " + f"(3 + 4); a cache reporting 0 is silently discarded by the model" + ) + # and the tensors really are concatenated, not just the bookkeeping patched + if hasattr(merged, "layers"): + assert merged.layers[0].keys.shape[-2] == 7 + assert merged.layers[0].values.shape[-2] == 7