From 50012d153d680a8ffd682fbb72f6069da5ce6c2f Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Tue, 23 Sep 2025 20:56:47 -0700 Subject: [PATCH] Add batch support for sliding window cache (#487) * add batch support for sliding window cache * fix * fix masks * fix cache * another test * comment --- mlx_lm/generate.py | 8 + mlx_lm/models/cache.py | 278 ++++++++++++++++++++++++++++++++--- mlx_lm/models/gemma3_text.py | 7 +- tests/test_generate.py | 51 +++++++ tests/test_models.py | 6 +- tests/test_prompt_cache.py | 104 ++++++++++++- 6 files changed, 418 insertions(+), 36 deletions(-) diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 5e01204..cd07287 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -26,8 +26,10 @@ from .models import cache from .models.cache import ( ArraysCache, BatchKVCache, + BatchRotatingKVCache, KVCache, QuantizedKVCache, + RotatingKVCache, load_prompt_cache, ) from .sample_utils import make_sampler @@ -866,6 +868,12 @@ def _make_cache(model, left_padding): elif isinstance(c, ArraysCache): c.left_padding = mx.array(left_padding) batch_cache.append(c) + elif isinstance(c, RotatingKVCache): + if c.keep > 0: + raise ValueError( + "RotatingKVCache with keep tokens is not supported." + ) + batch_cache.append(BatchRotatingKVCache(c.max_size, left_padding)) else: raise ValueError(f"{type(c)} does not yet support batching") return batch_cache diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index a56fecd..7d972dc 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -1,5 +1,6 @@ # Copyright © 2023-2024 Apple Inc. +import copy from typing import Any, Dict, List, Optional import mlx.core as mx @@ -73,10 +74,10 @@ def load_prompt_cache(file_name, return_metadata=False): arrays = tree_unflatten(list(arrays.items())) cache_metadata = tree_unflatten(list(cache_metadata.items())) info, metadata, classes = cache_metadata - cache = [globals()[c]() for c in classes] - for c, state, meta_state in zip(cache, arrays, info): - c.state = state - c.meta_state = meta_state + cache = [ + globals()[c].from_state(state, meta_state) + for c, state, meta_state in zip(classes, arrays, info) + ] if return_metadata: return cache, metadata return cache @@ -141,6 +142,14 @@ class _BaseCache: def is_trimmable(self): return False + @classmethod + def from_state(cls, state, meta_state): + # Create an instance of cls without calling __init__ + obj = cls.__new__(cls) + obj.state = state + obj.meta_state = meta_state + return obj + class ConcatenateKVCache(_BaseCache): """ConcatenateKVCache the simplest KV cache implementation. @@ -188,11 +197,12 @@ class ConcatenateKVCache(_BaseCache): class QuantizedKVCache(_BaseCache): + step = 256 + def __init__(self, group_size: int = 64, bits: int = 8): self.keys = None self.values = None self.offset = 0 - self.step = 256 self.group_size = group_size self.bits = bits @@ -254,11 +264,11 @@ class QuantizedKVCache(_BaseCache): @property def meta_state(self): - return tuple(map(str, (self.step, self.offset, self.group_size, self.bits))) + return tuple(map(str, (self.offset, self.group_size, self.bits))) @meta_state.setter def meta_state(self, v): - self.step, self.offset, self.group_size, self.bits = map(int, v) + self.offset, self.group_size, self.bits = map(int, v) def is_trimmable(self): return True @@ -273,11 +283,12 @@ class QuantizedKVCache(_BaseCache): class KVCache(_BaseCache): + step = 256 + def __init__(self): self.keys = None self.values = None self.offset = 0 - self.step = 256 def update_and_fetch(self, keys, values): prev = self.offset @@ -341,14 +352,14 @@ class KVCache(_BaseCache): class RotatingKVCache(_BaseCache): + step = 256 - def __init__(self, max_size=None, keep=0, step=256): + def __init__(self, max_size, keep=0): self.keep = keep self.keys = None self.values = None self.offset = 0 self.max_size = max_size - self.step = step self._idx = 0 def _trim(self, trim_size, v, append=None): @@ -389,9 +400,9 @@ class RotatingKVCache(_BaseCache): self.keys = self._temporal_order(self.keys) self.values = self._temporal_order(self.values) - # The largest size is self.max_size + S to ensure + # The largest size is self.max_size + S - 1 to ensure # every token gets at least self.max_size context - trim_size = self._idx - self.max_size + trim_size = self._idx - self.max_size + 1 self.keys = self._trim(trim_size, self.keys, keys) self.values = self._trim(trim_size, self.values, values) self.offset += keys.shape[2] @@ -459,13 +470,11 @@ class RotatingKVCache(_BaseCache): @property def meta_state(self): - return tuple( - map(str, (self.keep, self.max_size, self.step, self.offset, self._idx)) - ) + return tuple(map(str, (self.keep, self.max_size, self.offset, self._idx))) @meta_state.setter def meta_state(self, v): - self.keep, self.max_size, self.step, self.offset, self._idx = map( + self.keep, self.max_size, self.offset, self._idx = map( int, v, ) @@ -487,7 +496,7 @@ class RotatingKVCache(_BaseCache): ): if N > 1: window_size = window_size or self.max_size - offset = min(self.max_size, self.offset) + offset = min(self.max_size - 1, self.offset) if offset + N > window_size or return_array: return create_causal_mask(N, offset, window_size=window_size) else: @@ -500,16 +509,19 @@ class RotatingKVCache(_BaseCache): idx = self._idx if idx >= self.max_size: idx = 0 - mask_size = min(self.max_size, self.offset) + if self.offset < self.max_size: + mask_size = self.offset + 1 + else: + mask_size = self.max_size mask = mx.arange(mask_size) >= (mask_size - window_size) mask = mx.roll(mask, shift=idx + 1) - return mask[:, None] + return mask class ArraysCache(_BaseCache): def __init__(self, size, left_padding: Optional[List[int]] = None): self.cache = [None] * size - self.left_padding = left_padding + self.left_padding = mx.array(left_padding) if left_padding else None def __setitem__(self, idx, value): self.cache[idx] = value @@ -552,7 +564,7 @@ class MambaCache(ArraysCache): class ChunkedKVCache(KVCache): - def __init__(self, chunk_size=None): + def __init__(self, chunk_size): super().__init__() self.chunk_size = chunk_size self.start_position = 0 @@ -633,6 +645,8 @@ class CacheList(KVCache): class BatchKVCache(_BaseCache): + step = 256 + def __init__(self, left_padding: List[int]): """ The BatchKV cache expects inputs to be left-padded. @@ -657,7 +671,6 @@ class BatchKVCache(_BaseCache): self.left_padding = mx.array(left_padding) self.offset = mx.array([-l for l in left_padding]) self._idx = 0 - self.step = 256 def update_and_fetch(self, keys, values): prev = self._idx @@ -756,3 +769,224 @@ class BatchKVCache(_BaseCache): mx.concatenate, zip(*(pad(self), pad(other))) ) self._idx = max_idx + + +class BatchRotatingKVCache(_BaseCache): + step = 256 + + def __init__(self, max_size, left_padding: List[int]): + self.keys = None + self.values = None + + self.left_padding = mx.array(left_padding) + self.offset = mx.array([-l for l in left_padding]) + + self.max_size = max_size + self._idx = 0 + self._offset = 0 + self.rotated = False + + def _trim(self, trim_size, v, append=None): + if trim_size > 0: + v = v[..., trim_size:, :] + if append is not None: + return mx.concatenate([v, append], axis=2) + return v + + def _temporal_order(self): + """ + Rearrange the cache into temporal order. + """ + if self.rotated: + self.keys = mx.roll(self.keys, -self._idx, axis=2) + self.values = mx.roll(self.values, -self._idx, axis=2) + self._idx = self.keys.shape[2] + self.rotated = False + + def _update_concat(self, keys, values): + if self.keys is None: + self.keys = keys + self.values = values + else: + # Put the keys/values in temporal order to + # preserve context + self._temporal_order() + + # Slice off the end if needed + if self.keys.shape[2] > self._idx: + self.keys = self.keys[..., : self._idx, :] + self.values = self.values[..., : self._idx, :] + + # The largest size is self.max_size + S - 1 to ensure + # every token gets at least self.max_size context + trim_size = self._idx - self.max_size + 1 + if trim_size > 0: + self.left_padding -= trim_size + self.keys = self._trim(trim_size, self.keys, keys) + self.values = self._trim(trim_size, self.values, values) + self.offset += keys.shape[2] + self._offset += keys.shape[2] + self._idx = self.keys.shape[2] + return self.keys, self.values + + def _update_in_place(self, keys, values): + # May not have hit the max size yet, so potentially + # keep growing the cache + B, n_kv_heads, S, k_head_dim = keys.shape + prev = self._offset + if self.keys is None or ( + prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size + ): + v_head_dim = values.shape[3] + new_size = min(self.step, self.max_size - prev) + k_shape = (B, n_kv_heads, new_size, k_head_dim) + v_shape = (B, n_kv_heads, new_size, 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: + 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._idx = prev + + # Trim if needed + trim_size = self.keys.shape[2] - self.max_size + if trim_size > 0: + self.keys = self._trim(trim_size, self.keys) + self.values = self._trim(trim_size, self.values) + self._idx = self.max_size + self.left_padding -= trim_size + + # Rotate + if self._idx == self.max_size: + self.rotated = True + self._idx = 0 + if self.rotated: + self.left_padding -= S + + # Assign + self.keys[..., self._idx : self._idx + S, :] = keys + self.values[..., self._idx : self._idx + S, :] = values + self._offset += S + self.offset += S + self._idx += S + + # If the buffer is not full, slice off the end + if self._offset < self.max_size: + return ( + self.keys[..., : self._offset, :], + self.values[..., : self._offset, :], + ) + return self.keys, self.values + + def update_and_fetch(self, keys, values): + if keys.shape[2] == 1: + return self._update_in_place(keys, values) + return self._update_concat(keys, values) + + @property + def state(self): + k, v = self.keys, self.values + if self._offset < k.shape[2]: + k, v = k[..., : self._offset, :], v[..., : self._offset, :] + return k, v, self.offset, self.left_padding + + @state.setter + def state(self, v): + self.keys, self.values, self.offset, self.left_padding = v + + @property + def meta_state(self): + return tuple(map(str, (self.max_size, self._offset, self._idx, self.rotated))) + + @meta_state.setter + def meta_state(self, v): + self.max_size, self._offset, self._idx = map( + int, + v[:3], + ) + self.rotated = bool(v[3]) + + def is_trimmable(self): + return self._offset < self.max_size + + def trim(self, n): + n = min(self._offset, n) + self._offset -= n + self._idx -= n + self.offset -= n + return n + + def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache: + raise NotImplementedError("BatchRotatingKVCache Quantization NYI") + + def make_mask( + self, N: int, window_size: Optional[int] = None, return_array: bool = False + ): + left_padding = self.left_padding + window_size = window_size or self.max_size + offset = min(self.max_size - 1, self._offset) + rinds = mx.arange(offset + N) + linds = mx.arange(offset, offset + N) if offset else rinds + linds = linds[:, None] + rinds = rinds[None] + mask = linds >= rinds + mask &= linds < rinds + window_size + if (trim_size := self._idx - self.max_size + int(N > 1)) > 0: + left_padding = left_padding - trim_size + + rotated = N == 1 and (self.rotated or self._idx >= self.max_size) + if rotated: + left_padding = left_padding - 1 + + mask = mask & (rinds >= mx.expand_dims(left_padding, (1, 2, 3))) + + if rotated: + idx = self._idx + if idx >= self.max_size: + idx = 0 + mask = mx.roll(mask, shift=idx + 1, axis=-1) + + return mask + + def filter(self, batch_indices): + """ + In-place filter to keep just the given indices in the cache. + """ + self.keys = self.keys[batch_indices] + self.values = self.values[batch_indices] + self.offset = self.offset[batch_indices] + self.left_padding = self.left_padding[batch_indices] + + def extend(self, other): + """ + In-place extend this cache with the other cache. + """ + if (self.rotated != other.rotated) or self._idx != other._idx: + self._temporal_order() + other._temporal_order() + + max_idx = max(self._idx, other._idx) + max_size = max(self.keys.shape[2], other.keys.shape[2]) + + def pad(c): + left = max_idx - c._idx + right = max_size - c.keys.shape[2] - left + k, v = c.keys, c.values + if right < 0: + k = k[..., :right, :] + v = v[..., :right, :] + right = 0 + if left != 0 or right != 0: + pad = [(0, 0), (0, 0), (left, right), (0, 0)] + k = mx.pad(k, pad) + v = mx.pad(v, pad) + left_padding = c.left_padding + left + return k, v, c.offset, left_padding + + self.keys, self.values, self.offset, self.left_padding = map( + mx.concatenate, zip(*(pad(self), pad(other))) + ) + self._idx = max_idx + self._offset = max(self._offset, other._offset) diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py index 97531e9..f8e60d8 100644 --- a/mlx_lm/models/gemma3_text.py +++ b/mlx_lm/models/gemma3_text.py @@ -87,8 +87,6 @@ class Attention(nn.Module): keys = self.rope(keys) # Sliding window - if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]: - mask = mask[..., -keys.shape[-2] :] output = scaled_dot_product_attention( queries, keys, values, cache=cache, scale=self.scale, mask=mask ) @@ -194,7 +192,6 @@ class Gemma3Model(nn.Module): cache[0], window_size=self.window_size, ) - for i, (layer, c) in enumerate(zip(self.layers, cache)): is_global = ( i % self.sliding_window_pattern == self.sliding_window_pattern - 1 @@ -246,7 +243,5 @@ class Model(nn.Module): ): caches.append(KVCache()) else: - caches.append( - RotatingKVCache(max_size=self.args.sliding_window, keep=0) - ) + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) return caches diff --git a/tests/test_generate.py b/tests/test_generate.py index 628ec5d..0fd7e03 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -11,6 +11,7 @@ from mlx_lm.generate import ( generate, stream_generate, ) +from mlx_lm.models.cache import RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler from mlx_lm.utils import load @@ -301,6 +302,56 @@ class TestGenerate(unittest.TestCase): batch_tokens = batch_responses[uids[e]] self.assertEqual(tokens, batch_tokens) + def test_batch_sliding_window(self): + prompts = [ + "Write a story about Einstein", + "Hi", + "What time is it?", + "How tall is Mt Everest?", + ] + prompts = [ + self.tokenizer.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=True, + add_generation_prompt=True, + ) + for p in prompts + ] + + self.model.make_cache = lambda: [ + RotatingKVCache(max_size=4) for _ in self.model.layers + ] + batch_gen = BatchGenerator( + self.model, + stop_tokens=self.tokenizer.eos_token_ids, + max_tokens=10, + prefill_batch_size=1, + prefill_step_size=8, + completion_batch_size=2, + ) + uids = batch_gen.insert(prompts) + batch_responses = {uid: [] for uid in uids} + while responses := batch_gen.next(): + for r in responses: + batch_responses[r.uid].append(r.logprobs) + + for e, uid in enumerate(uids): + for i, response in enumerate( + stream_generate( + self.model, + self.tokenizer, + prompts[e], + max_tokens=10, + ) + ): + batch_logprobs = batch_responses[uid][i] + logprobs = response.logprobs + self.assertTrue( + mx.allclose(batch_logprobs, logprobs, rtol=1e-4, atol=1e-4) + ) + + del self.model.make_cache + if __name__ == "__main__": unittest.main() diff --git a/tests/test_models.py b/tests/test_models.py index b475d3b..4950f12 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -37,7 +37,7 @@ class TestModels(unittest.TestCase): def test_rotating_kv_cache(self): b, h, d = 1, 2, 32 - cache = RotatingKVCache(max_size=8, step=4) + cache = RotatingKVCache(max_size=8) k = mx.random.uniform(shape=(b, h, 2, d)) v = mx.random.uniform(shape=(b, h, 2, d)) @@ -70,7 +70,7 @@ class TestModels(unittest.TestCase): idx %= 8 # Try with nonzero keep - cache = RotatingKVCache(max_size=8, step=4, keep=2) + cache = RotatingKVCache(max_size=8, keep=2) # Check a large update k = mx.random.uniform(shape=(b, h, 20, d)) @@ -98,7 +98,7 @@ class TestModels(unittest.TestCase): # alternating prompt/prefill with generation d = 4 h = 2 - cache = RotatingKVCache(max_size=18, step=4) + cache = RotatingKVCache(max_size=18) x = mx.random.uniform(shape=(1, h, 8, d)) k, v = cache.update_and_fetch(x, x) diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py index 8dd3725..9b169ec 100644 --- a/tests/test_prompt_cache.py +++ b/tests/test_prompt_cache.py @@ -11,6 +11,7 @@ from mlx_lm.generate import generate_step from mlx_lm.models.base import create_attention_mask, create_causal_mask from mlx_lm.models.cache import ( BatchKVCache, + BatchRotatingKVCache, CacheList, ChunkedKVCache, KVCache, @@ -391,7 +392,7 @@ class TestPromptCache(unittest.TestCase): kv = mx.zeros((1, 1, 10, 32)) cache.update_and_fetch(kv, kv) mask = cache.make_mask(3, window_size=5) - self.assertEqual(mask.shape, (3, 11)) + self.assertEqual(mask.shape, (3, 10)) self.assertTrue(mx.all(mask.sum(axis=-1) == 5)) for i in range(3): s = 11 - 3 + i @@ -405,7 +406,7 @@ class TestPromptCache(unittest.TestCase): self.assertEqual(mask, None) mask = cache.make_mask(1, window_size=5) - self.assertEqual(mask.squeeze(1).tolist(), [True] + [False] * 3 + [True] * 4) + self.assertEqual(mask.tolist(), [True] + [False] * 3 + [True] * 4) cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5) self.assertTrue(mx.array_equal(cmask, mask)) @@ -413,9 +414,7 @@ class TestPromptCache(unittest.TestCase): cache.update_and_fetch(kv, kv) mask = cache.make_mask(1, window_size=5) - self.assertEqual( - mask.squeeze(1).tolist(), [True] * 2 + [False] * 3 + [True] * 3 - ) + self.assertEqual(mask.tolist(), [True] * 2 + [False] * 3 + [True] * 3) cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5) self.assertTrue(mx.array_equal(cmask, mask)) @@ -460,6 +459,101 @@ class TestPromptCache(unittest.TestCase): self.assertEqual(cache_a.offset.tolist(), [6, 7, 6, 1, 4]) self.assertEqual(cache_a.left_padding.tolist(), [2, 1, 2, 7, 4]) + def test_batch_rotating_kv_cache(self): + cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0]) + mask = cache.make_mask(4) + self.assertFalse(mx.any(mask[0, 0, 0, :])) + self.assertTrue( + mx.array_equal(mask[1, 0, 0, :], mx.array([True, False, False, False])) + ) + + # Batch update works + k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8)) + k, v = cache.update_and_fetch(k, v) + + mask = cache.make_mask(4) + k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8)) + k, v = cache.update_and_fetch(k, v) + self.assertEqual(mask.shape[-2:], (4, k.shape[2])) + self.assertEqual( + mask[0, 0, 0, :].tolist(), [False, True, True, True, False, False, False] + ) + + # Single query update works + cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0]) + k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8)) + k, v = cache.update_and_fetch(k, v) + + mask = cache.make_mask(1) + k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8)) + + k, v = cache.update_and_fetch(k, v) + self.assertEqual(mask.shape[-2:], (1, k.shape[2])) + self.assertEqual(mask[0, 0, 0].tolist(), [True, False, True, True]) + self.assertEqual(mask[1, 0, 0].tolist(), [True, True, True, True]) + + # Check filtering + cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0, 3]) + k, v = mx.zeros((3, 1, 3, 8)), mx.zeros((3, 1, 3, 8)) + cache.update_and_fetch(k, v) + cache.filter(mx.array([1])) + self.assertEqual(cache.keys.shape, (1, 1, 3, 8)) + + # Check extend + cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 1]) + other = BatchRotatingKVCache(max_size=4, left_padding=[2, 2]) + k, v = mx.zeros((2, 1, 5, 8)), mx.zeros((2, 1, 5, 8)) + cache.update_and_fetch(k, v) + other.update_and_fetch(k, v) + k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8)) + cache.update_and_fetch(k, v) + cache.extend(other) + + # Check mask when going from prompt -> extend -> prompt + cache = BatchRotatingKVCache(max_size=8, left_padding=[4]) + k, v = mx.zeros((1, 1, 8, 8)), mx.zeros((1, 1, 8, 8)) + cache.update_and_fetch(k, v) + + mask = cache.make_mask(1) + self.assertEqual( + mask.squeeze().tolist(), [True, False, False, False, True, True, True, True] + ) + + k, v = mx.zeros((1, 1, 1, 8)), mx.zeros((1, 1, 1, 8)) + cache.update_and_fetch(k, v) + + mask = cache.make_mask(2) + expected = mx.array( + [ + [False, False, False, True, True, True, True, True, False], + [False, False, False, True, True, True, True, True, True], + ] + ) + self.assertTrue(mx.array_equal(mask.squeeze(), expected)) + + def test_save_load_batch_caches(self): + cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors") + + cache = [ + MambaCache(left_padding=[1, 2]), + BatchKVCache(left_padding=[1, 2]), + BatchRotatingKVCache(max_size=10, left_padding=[1, 2]), + ] + for c in cache: + if isinstance(c, MambaCache): + c[0] = mx.random.uniform(shape=(4, 4, 4)) + c[1] = mx.random.uniform(shape=(4, 4, 4)) + else: + x = mx.random.uniform(shape=(4, 4, 7, 4)) + y = mx.random.uniform(shape=(4, 4, 7, 4)) + c.update_and_fetch(x, y) + + save_prompt_cache(cache_file, cache) + loaded_cache = load_prompt_cache(cache_file) + left_padding = mx.array([1, 2]) + for c, lc in zip(cache, loaded_cache): + self.assertTrue(mx.array_equal(c.left_padding, left_padding)) + if __name__ == "__main__": unittest.main()