fix: fix bug when calling _concat_caches in kv.py (from pr#177)

fix bug when calling _concat_caches in kv.py from pr#177 by @RixinLiu.

Author: spitzblattr

src/memos/memories/activation/kv.py

@@ -210,14 +210,15 @@ def _concat_caches(self, caches: list[DynamicCache]) -> DynamicCache:
 
         merged = DynamicCache()
         num_layers = len(caches[0].key_cache)
+        merged.append_new_layers(num_layers - 1)
 
         for layer in range(num_layers):
             # gather all K and V for this layer
-            keys = [c.key_cache[layer] for c in caches]
-            vals = [c.value_cache[layer] for c in caches]
+            keys = [c.layers[layer].keys for c in caches]
+            vals = [c.layers[layer].values for c in caches]
             # single concat per layer
-            merged.key_cache.append(torch.cat(keys, dim=-2))
-            merged.value_cache.append(torch.cat(vals, dim=-2))
+            merged.layers[layer].keys = (torch.cat(keys, dim=-2))
+            merged.layers[layer].values = (torch.cat(vals, dim=-2))
 
         return merged