Fix ArraysCache.from_state not initializing left_padding and lengths (ml-explore#807)

* Fix ArraysCache.from_state not initializing left_padding and lengths

The base class from_state uses __new__ to bypass __init__, but
ArraysCache.state.setter only sets self.cache, not left_padding
or lengths. This causes AttributeError when using loaded caches
with models that use MambaCache (e.g., Qwen3-Next).

The fix overrides from_state in ArraysCache to properly initialize
these attributes before setting state.

* Add test

---------

Co-authored-by: Awni Hannun <awni@apple.com>

Author: lpalbou

mlx_lm/models/cache.py

@@ -561,10 +561,16 @@ def empty(self):
 
 
 class ArraysCache(_BaseCache):
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+        instance.left_padding = None
+        instance.lengths = None
+        return instance
+
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = mx.array(left_padding) if left_padding else None
-        self.lengths = None
+        if left_padding:
+            self.left_padding = mx.array(left_padding)
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value

tests/test_prompt_cache.py

@@ -133,6 +133,19 @@ 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_mamba_cache(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [MambaCache()]
+        cache[0][0] = mx.zeros((1, 4, 4))
+        cache[0][1] = mx.zeros((1, 4, 4))
+
+        save_prompt_cache(cache_file, cache)
+        loaded = load_prompt_cache(cache_file)
+
+        # Try to make a mask
+        mask = loaded[0].make_mask(4)
+
     def test_cache_with_generate(self):
         model, tokenizer = self.model, self.tokenizer
         prompt = tokenizer.encode("this is a prompt", return_tensors="mlx")[0]