1"""Cache adapters that manipulate real mlx_lm (0.32.x) KV tensors.23Dense mlx_lm KVCache stores post-RoPE keys/values [B, H_kv, T_buf, d_h] with buffer growth in steps of 256 and a4logical `offset`. QuantizedKVCache stores packed uint32 words + scales + biases per group (mx.quantize) and is consumed5by quantized_scaled_dot_product_attention (genuine quantized execution).67SparseKVCache / SparseQuantizedKVCache: keep a subset of positions (shared or per kv-head gather) while preserving8the *logical* offset so newly appended tokens get their original RoPE positions; retained keys already carry their9own positions (RoPE applied before caching). Multi-token appends use an explicit boolean mask: retained positions are10fully visible, new tokens are causal among themselves.11"""12from __future__ import annotations13from typing import Optional14import mlx.core as mx15from mlx.utils import tree_map16from mlx_lm.models.cache import KVCache, QuantizedKVCache171819def dense_kv(c) -> tuple[mx.array, mx.array]:20 """Trimmed (offset-length) dense keys/values of a dense KVCache."""21 return c.keys_and_values()222324def _bool_mask(n_new: int, n_kept: int) -> mx.array:25 left = mx.ones((n_new, n_kept), dtype=mx.bool_)26 r = mx.arange(n_new)27 right = r[:, None] >= r[None, :]28 return mx.concatenate([left, right], axis=1)293031class SparseKVCache(KVCache):32 """KV cache holding a subset of positions with a preserved logical offset (dense dtype)."""3334 def __init__(self):35 super().__init__()36 self.n_kept = 037 self.retained_positions = None3839 @classmethod40 def from_dense(cls, keys: mx.array, values: mx.array, logical_offset: int, indices=None) -> "SparseKVCache":41 c = cls()42 if indices is None:43 nk, nv = keys, values44 kept = int(keys.shape[2])45 else:46 idx = mx.array(indices, dtype=mx.int32)47 if idx.ndim == 1:48 nk = mx.take(keys, idx, axis=2); nv = mx.take(values, idx, axis=2); kept = int(idx.shape[0])49 elif idx.ndim == 2:50 g = idx[None, :, :, None]51 nk = mx.take_along_axis(keys, mx.broadcast_to(g, (keys.shape[0], keys.shape[1], idx.shape[1], keys.shape[3])), axis=2)52 nv = mx.take_along_axis(values, mx.broadcast_to(g, (values.shape[0], values.shape[1], idx.shape[1], values.shape[3])), axis=2)53 kept = int(idx.shape[1])54 else:55 raise ValueError("indices must be 1-D or 2-D")56 c.retained_positions = idx.tolist()57 c.keys, c.values = nk, nv58 c.n_kept = kept59 c.offset = int(logical_offset)60 mx.eval(c.keys, c.values)61 return c6263 @classmethod64 def from_kvcache(cls, src: KVCache, indices=None, logical_offset: Optional[int] = None) -> "SparseKVCache":65 k, v = dense_kv(src)66 return cls.from_dense(k, v, src.offset if logical_offset is None else logical_offset, indices)6768 def update_and_fetch(self, keys, values):69 n_new = keys.shape[2]70 if self.keys is None:71 self.keys, self.values = keys, values72 else:73 self.keys = mx.concatenate([self.keys[..., : self.n_kept, :], keys], axis=2)74 self.values = mx.concatenate([self.values[..., : self.n_kept, :], values], axis=2)75 self.n_kept += n_new76 self.offset += n_new77 return self.keys, self.values7879 def keys_and_values(self):80 return self.keys[..., : self.n_kept, :], self.values[..., : self.n_kept, :]8182 @property83 def state(self):84 return self.keys, self.values, self.offset, self.n_kept8586 @state.setter87 def state(self, v):88 self.keys, self.values, self.offset, self.n_kept = v8990 def size(self):91 return self.n_kept9293 def make_mask(self, N, return_array=False, window_size=None):94 if N == 1:95 return None96 return _bool_mask(N, self.n_kept)9798 def is_trimmable(self):99 return False100101 @property102 def nbytes(self):103 if self.keys is None:104 return 0105 k, v = self.keys_and_values()106 return int(k.nbytes + v.nbytes)107108 def to_quantized(self, group_size: int = 64, bits: int = 4) -> "SparseQuantizedKVCache":109 q = SparseQuantizedKVCache(group_size=group_size, bits=bits)110 q.offset = self.offset111 q.n_kept = self.n_kept112 k, v = self.keys_and_values()113 q.keys = mx.quantize(k, group_size=group_size, bits=bits)114 q.values = mx.quantize(v, group_size=group_size, bits=bits)115 mx.eval(q.keys, q.values)116 return q117118119class SparseQuantizedKVCache(QuantizedKVCache):120 """Packed quantized cache over retained positions with preserved logical offset."""121122 def __init__(self, group_size: int = 64, bits: int = 8):123 super().__init__(group_size=group_size, bits=bits)124 self.n_kept = 0125126 def update_and_fetch(self, keys, values):127 n_new = keys.shape[2]128 qk = mx.quantize(keys, group_size=self.group_size, bits=self.bits)129 qv = mx.quantize(values, group_size=self.group_size, bits=self.bits)130 if self.keys is None:131 self.keys, self.values = qk, qv132 else:133 cur_k, cur_v = self.keys_and_values()134 self.keys = tuple(mx.concatenate([a, b], axis=2) for a, b in zip(cur_k, qk))135 self.values = tuple(mx.concatenate([a, b], axis=2) for a, b in zip(cur_v, qv))136 self.n_kept += n_new137 self.offset += n_new138 return self.keys, self.values139140 def keys_and_values(self):141 return (tree_map(lambda x: x[..., : self.n_kept, :], self.keys),142 tree_map(lambda x: x[..., : self.n_kept, :], self.values))143144 @property145 def state(self):146 return self.keys, self.values, self.offset, self.group_size, self.bits, self.n_kept147148 @state.setter149 def state(self, v):150 self.keys, self.values, self.offset, self.group_size, self.bits, self.n_kept = v151152 def size(self):153 return self.n_kept154155 def make_mask(self, N, return_array=False, window_size=None):156 if N == 1:157 return None158 return _bool_mask(N, self.n_kept)159160 def is_trimmable(self):161 return False162163 @property164 def nbytes(self):165 if self.keys is None:166 return 0167 k, v = self.keys_and_values()168 return int(sum(x.nbytes for x in k) + sum(x.nbytes for x in v))169170171# ---------------------------------------------------------------------------------------------------------------172def to_sparse_all(cache: list) -> list:173 """Wrap every dense layer into a SparseKVCache retaining all positions (100% retention control / B0 canonical form)."""174 return [SparseKVCache.from_kvcache(c) if type(c) is KVCache else c for c in cache]175176177def quantize_cache(cache: list, bits: int, group_size: int = 64) -> list:178 """B2: packed quantization of every layer's K/V (storage AND execution representation)."""179 out = []180 for c in cache:181 if isinstance(c, SparseKVCache):182 out.append(c.to_quantized(group_size=group_size, bits=bits))183 elif type(c) is KVCache:184 out.append(SparseKVCache.from_kvcache(c).to_quantized(group_size=group_size, bits=bits))185 else:186 raise TypeError(f"cannot quantize cache of type {type(c)}")187 return out188189190def select_positions(cache: list, indices_per_layer, logical_offset: Optional[int] = None) -> list:191 """B3/B4/B5/B6: sparse cache retaining the given positions per layer (1-D shared or 2-D per kv-head indices)."""192 out = []193 for c, idx in zip(cache, indices_per_layer):194 if isinstance(c, SparseKVCache):195 k, v = c.keys_and_values()196 out.append(SparseKVCache.from_dense(k, v, c.offset if logical_offset is None else logical_offset, idx))197 elif type(c) is KVCache:198 out.append(SparseKVCache.from_kvcache(c, idx, logical_offset))199 else:200 raise TypeError("select_positions expects dense layers")201 return out202203204def copy_cache(cache: list) -> list:205 """Independent branch of a cache (no shared mutable arrays)."""206 out = []207 for c in cache:208 if isinstance(c, SparseQuantizedKVCache):209 n = SparseQuantizedKVCache(group_size=c.group_size, bits=c.bits)210 k, v = c.keys_and_values()211 n.keys = tuple(mx.array(x) for x in k); n.values = tuple(mx.array(x) for x in v)212 n.offset, n.n_kept = c.offset, c.n_kept213 elif isinstance(c, SparseKVCache):214 k, v = c.keys_and_values()215 n = SparseKVCache(); n.keys, n.values = mx.array(k), mx.array(v); n.offset, n.n_kept = c.offset, c.n_kept216 n.retained_positions = c.retained_positions217 elif type(c) is KVCache:218 k, v = c.keys_and_values()219 n = KVCache(); n.keys, n.values = mx.array(k), mx.array(v); n.offset = c.offset220 elif type(c) is QuantizedKVCache:221 k, v = c.keys_and_values()222 n = QuantizedKVCache(group_size=c.group_size, bits=c.bits)223 n.keys = tuple(mx.array(x) for x in k); n.values = tuple(mx.array(x) for x in v); n.offset = c.offset224 else:225 raise TypeError(type(c))226 out.append(n)227 mx.eval([x for c in out for x in _arrays_of(c)])228 return out229230231def _arrays_of(c) -> list[mx.array]:232 k, v = c.keys_and_values()233 if isinstance(k, tuple):234 return list(k) + list(v)235 return [k, v]236237238def cache_tensor_bytes(cache: list) -> int:239 """Stored tensor bytes over the *logical content* (trimmed to stored positions, not allocator buffers)."""240 return int(sum(c.nbytes if not (type(c) is KVCache or type(c) is QuantizedKVCache) else sum(x.nbytes for x in _arrays_of(c)) for c in cache))241242243def dense_bytes_formula(n_layers: int, T: int, n_kv_heads: int, head_dim: int, bytes_per_scalar: int, B: int = 1) -> int:244 """M_KV = 2 B L T H_kv d_h s (consistency check only; measured bytes are authoritative)."""245 return 2 * B * n_layers * T * n_kv_heads * head_dim * bytes_per_scalar246