Fix sliding window batching (ml-explore#738)

Author: Awni Hannun

mlx_lm/generate.py

@@ -1071,7 +1071,8 @@ def _process_prompts(self, prompts):
             prompt_cache = _merge_caches(caches)
 
             for c in prompt_cache:
-                c.prepare(lengths=lengths, right_padding=padding)
+                # subtract one from lengths since we don't process the last token during prefill
+                c.prepare(lengths=[l - 1 for l in lengths], right_padding=padding)
 
             while inputs.shape[1] > 1:
                 n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
@@ -1096,6 +1097,7 @@ def _process_prompts(self, prompts):
         y, logprobs = self._step(
             inputs, prompt_cache, samplers, logits_processors, tokens
         )
+
         mx.async_eval(y, logprobs)
 
         return Batch(

mlx_lm/models/cache.py

@@ -1157,12 +1157,10 @@ def extract(self, idx):
             cache.keys = mx.roll(cache.keys, -self._idx, axis=2)
             cache.values = mx.roll(cache.values, -self._idx, axis=2)
             cache._idx = self.max_size
-        if padding > 0:
-            cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
-            cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
+        cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
+        cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
         cache.offset = offset
         cache._idx = cache.keys.shape[2]
-
         return cache
 
     @classmethod
@@ -1185,8 +1183,8 @@ def merge(cls, caches):
         keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
         values = mx.zeros((B, H, max_length, Dv), dtype=dt)
         for i, (p, c) in enumerate(zip(padding, caches)):
-            keys[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.keys)
-            values[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.values)
+            keys[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.keys)
+            values[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.values)
 
         cache = cls(caches[0].max_size, padding)
         cache.keys = keys

tests/test_generate.py

@@ -473,7 +473,7 @@ def test_batch_continued_generation(self):
                 self.model,
                 stop_tokens=self.tokenizer.eos_token_ids,
                 max_tokens=10,
-                prefill_batch_size=1,
+                prefill_batch_size=4,
                 prefill_step_size=8,
                 completion_batch_size=2,
             )