Bring back max-kv-size to the batch generator (ml-explore#1106)

Author: angeloskath

mlx_lm/generate.py

@@ -1498,6 +1498,7 @@ def __init__(
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        max_kv_size: Optional[int] = None,
     ):
         self.model = model
         self.max_tokens = max_tokens
@@ -1507,6 +1508,7 @@ def __init__(
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.max_kv_size = max_kv_size
 
         self._default_state_machine = SequenceStateMachine(
             {"normal": [(seq, None) for seq in stop_tokens]} if stop_tokens else {},
@@ -1613,7 +1615,7 @@ def insert_segments(
         caches = caches or [None] * len(segments)
         for i in range(len(segments)):
             if caches[i] is None:
-                caches[i] = cache.make_prompt_cache(self.model)
+                caches[i] = self._make_new_cache()
 
         for seq, m, c, at, s, lp, sm in zip(
             segments,
@@ -1636,6 +1638,19 @@ def insert_segments(
 
         return uids
 
+    def _make_new_cache(self):
+        if self.max_kv_size is None:
+            return cache.make_prompt_cache(self.model)
+
+        return [
+            (
+                RotatingKVCache(max_size=self.max_kv_size)
+                if isinstance(ci, KVCache)
+                else ci
+            )
+            for ci in cache.make_prompt_cache(self.model)
+        ]
+
     def _find_uids(self, uids):
         uids = set(uids)
         results = {}

mlx_lm/server.py

@@ -1345,12 +1345,13 @@ def generate_response(
         # Add dynamic response
         if self.object_type.startswith("chat.completion"):
             key_name = "delta" if self.stream else "message"
-            choice[key_name] = {
-                "role": "assistant",
-                "content": text,
-                "reasoning": reasoning_text,
-                "tool_calls": tool_calls,
-            }
+            choice[key_name] = {"role": "assistant"}
+            if text:
+                choice[key_name]["content"] = text
+            if reasoning_text:
+                choice[key_name]["reasoning"] = reasoning_text
+            if tool_calls:
+                choice[key_name]["tool_calls"] = tool_calls
         elif self.object_type == "text_completion":
             choice.update(text=text)
         else:

tests/test_generate.py

@@ -15,7 +15,7 @@
     generate_step,
     stream_generate,
 )
-from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.models.cache import KVCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.utils import load
 
@@ -719,6 +719,57 @@ def test_remove_prompt_batch_updates_currently_processing(self):
             if all(r.finish_reason is not None for r in responses):
                 break
 
+    def test_batch_max_kv_size_creates_rotating_cache(self):
+        max_kv_size = 256
+        gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            max_kv_size=max_kv_size,
+        )
+
+        prompt = self.tokenizer.encode("Write a long story about a cat")
+        gen.insert([prompt])
+
+        for r in gen.next_generated():
+            if r.finish_reason is not None:
+                for cache in r.prompt_cache:
+                    self.assertIsInstance(cache, RotatingKVCache)
+                    self.assertEqual(cache.max_size, max_kv_size)
+
+    def test_batch_max_kv_size_limits_cache_growth(self):
+        max_kv_size = 5
+        gen = BatchGenerator(
+            self.model,
+            max_tokens=10,
+            max_kv_size=max_kv_size,
+            prefill_batch_size=1,
+            prefill_step_size=128,
+            completion_batch_size=1,
+        )
+
+        prompt = self.tokenizer.encode("Write a long story about a cat")
+        gen.insert([prompt])
+
+        for r in gen.next_generated():
+            if r.finish_reason is not None:
+                for cache in r.prompt_cache:
+                    self.assertLessEqual(cache.keys.shape[2], max_kv_size)
+
+    def test_batch_max_kv_size_none_creates_regular_cache(self):
+        gen = BatchGenerator(
+            self.model,
+            max_tokens=1,
+            max_kv_size=None,
+        )
+
+        prompt = self.tokenizer.encode("Write a long story about a cat")
+        gen.insert([prompt])
+
+        for r in gen.next_generated():
+            if r.finish_reason is not None:
+                for cache in r.prompt_cache:
+                    self.assertIsInstance(cache, KVCache)
+
 
 if __name__ == "__main__":
     unittest.main()