Repository: ml-explore/mlx-lm

Issue:
KeyError: 'ChunkedKVCache' inside cache.py when restoring llama4 prompt cache
There's currently an issue with restoring state from a prompt cache file when using llama4.

```sh
cat document.txt | mlx_lm.cache_prompt --model Llama-4-Scout-17B-16E-4bit --prompt - --prompt-cache-file temp.safetensors
Processed   1540 tokens (152.62 tok/s)
mx.metal.get_peak_memory is deprecated and will be removed in a future version. Use mx.get_peak_memory instead.
Peak memory: 62.120 GB
Saving...
```

```sh
mlx_lm.generate --prompt-cache-file temp.safetensors --prompt "\nSummarize the above text."                                                                                                            Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/miniconda/base/envs/brennan/bin/mlx_lm.generate", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/opt/homebrew/Caskroom/miniconda/base/envs/brennan/lib/python3.12/site-packages/mlx_lm/generate.py", line 718, in main
    prompt_cache, metadata = load_prompt_cache(
                             ^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniconda/base/envs/brennan/lib/python3.12/site-packages/mlx_lm/models/cache.py", line 74, in load_prompt_cache
    cache = [globals()[c]() for c in classes]
             ~~~~~~~~~^^^
KeyError: 'ChunkedKVCache'
```

Versions used:
mlx 0.24.2
mlx-lm 0.22.4

Discussion:

Solution patch:
diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py
index 14026f0..7ce0ba2 100644
--- a/mlx_lm/models/cache.py
+++ b/mlx_lm/models/cache.py
@@ -436,3 +436,55 @@ class MambaCache(_BaseCache):
     @state.setter
     def state(self, v):
         self.cache = v
+
+
+class ChunkedKVCache(KVCache):
+    def __init__(self, chunk_size=None):
+        super().__init__()
+        self.chunk_size = chunk_size
+        self.start_position = 0
+
+    def maybe_trim_front(self):
+        # Maintain the cache below the chunk size
+        if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
+            self.keys = self.keys[..., self.chunk_size :, :]
+            self.values = self.values[..., self.chunk_size :, :]
+            self.start_position += self.chunk_size
+
+    def update_and_fetch(self, keys, values):
+        prev = self.offset - self.start_position
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
+            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        end = self.offset - self.start_position
+        self.keys[..., prev:end, :] = keys
+        self.values[..., prev:end, :] = values
+        return self.keys[..., :end, :], self.values[..., :end, :]
+
+    def trim(self, n):
+        n = min(self.offset - self.start_position, n)
+        self.offset -= n
+        return n
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.chunk_size, self.start_position)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.chunk_size, self.start_position = map(int, v)
diff --git a/mlx_lm/models/deepseek_v2.py b/mlx_lm/models/deepseek_v2.py
index 7a5bdeb..9964481 100644
--- a/mlx_lm/models/deepseek_v2.py
+++ b/mlx_lm/models/deepseek_v2.py
@@ -148,7 +148,7 @@ class DeepseekV2Attention(nn.Module):
             self.q_a_proj = nn.Linear(
                 self.hidden_size, self.q_lora_rank, bias=config.attention_bias
             )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
             self.q_b_proj = nn.Linear(
                 self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
             )
@@ -158,7 +158,7 @@ class DeepseekV2Attention(nn.Module):
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=config.attention_bias,
         )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
         self.kv_b_proj = nn.Linear(
             self.kv_lora_rank,
             self.num_heads
@@ -400,8 +400,6 @@ class DeepseekV2Model(nn.Module):
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        # Hack to avoid time-outs during prompt-processing
-        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
         if mask is None:
             mask = create_attention_mask(h, cache)
 
@@ -410,19 +408,17 @@ class DeepseekV2Model(nn.Module):
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
         for i in range(self.num_layers):
             h = self.layers[self.start_idx + i](h, mask, cache[i])
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
-            h = mx.distributed.send(
-                h, (pipeline_rank - 1) % pipeline_size, stream=dist_stream
-            )
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
diff --git a/mlx_lm/models/llama4.py b/mlx_lm/models/llama4.py
index 4c83f5d..f8f168c 100644
--- a/mlx_lm/models/llama4.py
+++ b/mlx_lm/models/llama4.py
@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import KVCache
+from .cache import ChunkedKVCache, KVCache
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -234,25 +234,25 @@ class LlamaModel(nn.Module):
         if cache is not None:
             for idx, c in enumerate(cache):
                 if (idx + 1) % 4 != 0:
-                    c.update()
+                    c.maybe_trim_front()
+            start = cache[0].start_position
+            offset = cache[0].offset
+        else:
+            start = 0
+            offset = 0
+        end = offset + h.shape[1]
+        linds = mx.arange(start, end)
+        rinds = mx.arange(offset, end)[:, None]
+        block_pos = mx.abs(
+            (linds // self.attention_chunk_size) - (rinds // self.attention_chunk_size)
+        )
+        token_pos = linds <= rinds
+        chunk_mask = (block_pos == 0) & token_pos
 
         if mask is None:
             mask = create_attention_mask(h, cache)
-            if cache is not None:
-                start = cache[0].start_position
-                offset = cache[0].offset
-            else:
-                start = 0
-                offset = 0
-            end = offset + h.shape[1]
-            linds = mx.arange(start, end)
-            rinds = mx.arange(offset, end)[:, None]
-            block_pos = mx.abs(
-                (linds // self.attention_chunk_size)
-                - (rinds // self.attention_chunk_size)
-            )
-            token_pos = linds <= rinds
-            chunk_mask = (block_pos == 0) & token_pos
+        else:
+            chunk_mask &= mask
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -288,50 +288,6 @@ class LanguageModel(nn.Module):
         return self.lm_head(out)
 
 
-class ChunkedKVCache(KVCache):
-    def __init__(self, chunk_size):
-        super().__init__()
-        self.chunk_size = chunk_size
-        self.start_position = 0
-
-    def update(self):
-        # Maintain the cache below the chunk size
-        if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
-            self.keys = self.keys[..., self.chunk_size :, :]
-            self.values = self.values[..., self.chunk_size :, :]
-            self.start_position += self.chunk_size
-
-    def update_and_fetch(self, keys, values):
-        prev = self.offset - self.start_position
-        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
-            B, n_kv_heads, _, k_head_dim = keys.shape
-            v_head_dim = values.shape[3]
-            n_steps = (self.step + keys.shape[2] - 1) // self.step
-            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
-            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
-            new_k = mx.zeros(k_shape, keys.dtype)
-            new_v = mx.zeros(v_shape, values.dtype)
-            if self.keys is not None:
-                if prev % self.step != 0:
-                    self.keys = self.keys[..., :prev, :]
-                    self.values = self.values[..., :prev, :]
-                self.keys = mx.concatenate([self.keys, new_k], axis=2)
-                self.values = mx.concatenate([self.values, new_v], axis=2)
-            else:
-                self.keys, self.values = new_k, new_v
-
-        self.offset += keys.shape[2]
-        end = self.offset - self.start_position
-        self.keys[..., prev:end, :] = keys
-        self.values[..., prev:end, :] = values
-        return self.keys[..., :end, :], self.values[..., :end, :]
-
-    def trim(self, n):
-        n = min(self.offset - self.start_position, n)
-        self.offset -= n
-        return n
-
-
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()


Test patch:
diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py
index d299cd0..a060784 100644
--- a/tests/test_prompt_cache.py
+++ b/tests/test_prompt_cache.py
@@ -9,6 +9,7 @@ import mlx.core as mx
 
 from mlx_lm.generate import generate_step
 from mlx_lm.models.cache import (
+    ChunkedKVCache,
     KVCache,
     MambaCache,
     QuantizedKVCache,
@@ -95,7 +96,13 @@ class TestPromptCache(unittest.TestCase):
     def test_save_load_mixed_cache(self):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 
-        cache = [MambaCache(), KVCache(), RotatingKVCache(8), MambaCache()]
+        cache = [
+            MambaCache(),
+            KVCache(),
+            RotatingKVCache(8),
+            MambaCache(),
+            ChunkedKVCache(256),
+        ]
         for c in cache:
             if isinstance(c, MambaCache):
                 c[0] = mx.random.uniform(shape=(4, 4, 4))