fix: stop merged KV caches being silently discarded on transformers >= 4.57

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"):

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