Fix save/load of CacheList (ml-explore#886)

Author: angeloskath

mlx_lm/models/cache.py

@@ -742,16 +742,24 @@ def trim(self, n):
 
     @property
     def state(self):
-        return [s for c in self.caches for s in c.state]
+        return [c.state for c in self.caches]
 
     @state.setter
     def state(self, v):
-        state_lens = [len(c.state) for c in self.caches]
-        start = 0
-        for c in self.caches:
-            l = len(c.state)
-            c.state = v[start : start + l]
-            start += l
+        for c, s in zip(self.caches, v):
+            c.state = s
+
+    @property
+    def meta_state(self):
+        return (
+            [type(c).__name__ for c in self.caches],
+            [c.meta_state for c in self.caches],
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        for c, m in zip(self.caches, v[1]):
+            c.meta_state = m
 
     def filter(self, batch_indices):
         """
@@ -793,6 +801,14 @@ def size(self):
     def empty(self):
         return self.caches[0].empty()
 
+    @classmethod
+    def from_state(cls, state, meta_state):
+        obj = cls.__new__(cls)
+        obj.caches = [
+            globals()[c].from_state(s, m) for s, c, m in zip(state, *meta_state)
+        ]
+        return obj
+
 
 def dynamic_roll(x, shifts, axis):
     n = x.shape[axis]

tests/test_prompt_cache.py

@@ -132,6 +132,41 @@ def test_save_load_mixed_cache(self):
                 self.assertTrue(mx.array_equal(k, lk))
                 self.assertTrue(mx.array_equal(v, lv))
 
+    def test_save_load_cache_list(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [
+            ArraysCache(size=2),
+            KVCache(),
+            RotatingKVCache(8),
+            ArraysCache(size=2),
+            ChunkedKVCache(256),
+        ]
+        for c in cache:
+            if isinstance(c, ArraysCache):
+                c[0] = mx.random.uniform(shape=(4, 4, 4))
+                c[1] = mx.random.uniform(shape=(4, 4, 4))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 7, 4))
+                y = mx.random.uniform(shape=(4, 4, 7, 4))
+                c.update_and_fetch(x, y)
+        cache = [CacheList(*cache)]
+
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        for c, lc in zip(cache[0].caches, loaded_cache[0].caches):
+            if isinstance(c, ArraysCache):
+                self.assertTrue(mx.array_equal(c[0], lc[0]))
+                self.assertTrue(mx.array_equal(c[1], lc[1]))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 1, 4))
+                y = mx.random.uniform(shape=(4, 4, 1, 4))
+                k, v = c.update_and_fetch(x, y)
+                lk, lv = lc.update_and_fetch(x, y)
+                self.assertEqual(c.offset, lc.offset)
+                self.assertTrue(mx.array_equal(k, lk))
+                self.assertTrue(mx.array_equal(v, lv))
+
     def test_save_load_arrays_cache(self):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")