fix bug when calling _concat_caches in kv.py

Author: RixinLiu

src/memos/memories/activation/kv.py

@@ -211,13 +211,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