fix (#503)

Author: awni

mlx_lm/models/cache.py

@@ -399,6 +399,7 @@ def _update_concat(self, keys, values):
             # preserve context
             self.keys = self._temporal_order(self.keys)
             self.values = self._temporal_order(self.values)
+            self._idx = self.keys.shape[2]
 
             # The largest size is self.max_size + S - 1 to ensure
             # every token gets at least self.max_size context

tests/test_prompt_cache.py

@@ -554,6 +554,20 @@ def test_save_load_batch_caches(self):
         for c, lc in zip(cache, loaded_cache):
             self.assertTrue(mx.array_equal(c.left_padding, left_padding))
 
+    def test_rotating_cache_updates(self):
+        cache = RotatingKVCache(max_size=8)
+        k = v = mx.zeros((1, 1, 10, 1))
+        cache.update_and_fetch(k, v)
+
+        for _ in range(3):
+            k = v = mx.zeros((1, 1, 1, 1))
+            cache.update_and_fetch(k, v)
+
+        k = v = mx.zeros((1, 1, 3, 1))
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(k.shape[2], 10)
+        self.assertEqual(v.shape[2], 10)
+
 
 if __name__ == "__main__":
     unittest.main()