1"""Position-selection policies over real caches.23- sink_recent (B3): StreamingLLM-inspired *retention pattern* (first `sinks` positions + most recent). Unlike StreamingLLM,4 retained keys keep their ORIGINAL positions (StreamingLLM re-rotates to positions-within-cache); labelled 'inspired'.5- random (B4): uniform random positions, seeded; matched budget.6- attention_informed (B5): faithful re-implementation of SnapKV's selection rule (observation window = last W prompt7 tokens; per kv-head attention of window queries over prefix keys, averaged over the window queries, 1-D pooled with8 kernel k (avg or max), top-k per kv-head; window positions always kept). GQA handled by averaging query heads in each9 kv group. Requires the window queries: we capture them with a forward hook (an extra pass over the window tokens on a10 *copy* of the cache), so optimized fused attention is NOT disabled for the main computation; the extra pass cost is11 recorded as compression time.12- pyramid (B6): PyramidKV-style layer budget allocation (arithmetic decrease from bottom to top layer, total budget13 preserved) using the attention_informed scores within each layer; labelled 'faithful_reimplementation' of the budget14 rule (verified against KVCache-Factory code: base=cap-w, min=base//beta, max=2*base-min, steps=(max-min)//(L-1)).15"""16from __future__ import annotations17import math18import numpy as np19import mlx.core as mx20import mlx.nn as nn21from .sparse import SparseKVCache, select_positions, copy_cache, dense_kv222324def sink_recent_indices(T: int, keep: int, sinks: int = 4) -> np.ndarray:25 keep = max(min(keep, T), 1)26 s = min(sinks, keep)27 rec = keep - s28 idx = np.concatenate([np.arange(s), np.arange(T - rec, T)]) if rec > 0 else np.arange(s)29 return np.unique(idx)303132def random_indices(T: int, keep: int, seed: int) -> np.ndarray:33 rng = np.random.default_rng(seed)34 keep = max(min(keep, T), 1)35 return np.sort(rng.choice(T, keep, replace=False))363738class _QueryCapture:39 """Context manager: temporarily patch the attention class __call__ to record post-RoPE queries per layer."""4041 def __init__(self, model):42 self.model = model43 inner = model.model if hasattr(model, "model") and hasattr(model.model, "layers") else model44 self.layers = inner.layers45 self.queries: list = [None] * len(self.layers)46 self.index_of = {id(l.self_attn): i for i, l in enumerate(self.layers)}47 self.cls = type(self.layers[0].self_attn)48 self.orig = self.cls.__call__4950 def __enter__(self):51 cap = self; orig = self.orig5253 def patched(self_attn, x, mask=None, cache=None):54 idx = cap.index_of.get(id(self_attn))55 if idx is not None:56 B, L, D = x.shape57 q = self_attn.q_proj(x).reshape(B, L, self_attn.n_heads, -1).transpose(0, 2, 1, 3)58 q = self_attn.rope(q, offset=cache.offset) if cache is not None else self_attn.rope(q)59 cap.queries[idx] = q60 return orig(self_attn, x, mask=mask, cache=cache)61 self.cls.__call__ = patched62 return self6364 def __exit__(self, *a):65 self.cls.__call__ = self.orig666768def attention_window_scores(model, cache: list, window_token_ids: list[int], pool_kernel: int = 7, pooling: str = "max") -> list[np.ndarray]:69 """Return per-layer scores [H_kv, T_prefix] = pooled mean attention of the window queries onto prefix keys.70 The window tokens are processed on a *copy* of the cache (the original stays immutable)."""71 if hasattr(model, "model") and not hasattr(model, "layers") and hasattr(model, "info"):72 model = model.model # MLXModel wrapper -> underlying nn.Module73 work = copy_cache(cache)74 T = work[0].offset75 toks = mx.array(window_token_ids, dtype=mx.int32)[None]76 with _QueryCapture(model) as qc:77 model(toks, cache=work)78 mx.eval([q for q in qc.queries])79 scores = []80 for layer_idx, c in enumerate(work):81 k, _ = c.keys_and_values() # [B, H_kv, T+W, d]82 k = k[..., :T, :].astype(mx.float32)83 q = qc.queries[layer_idx].astype(mx.float32) # [B, H_q, W, d]84 Hq, Hkv = q.shape[1], k.shape[1]85 rep = Hq // Hkv86 q = q.reshape(q.shape[0], Hkv, rep, q.shape[2], q.shape[3]) # [B, H_kv, rep, W, d]87 s = (q @ k[:, :, None, :, :].transpose(0, 1, 2, 4, 3)) / math.sqrt(k.shape[-1]) # [B, H_kv, rep, W, T]88 s = mx.softmax(s, axis=-1)89 s = s.mean(axis=(0, 2, 3)) # [H_kv, T]: mean over batch, group heads, window queries90 a = np.array(s)91 if pool_kernel > 1:92 pad = pool_kernel // 293 ap = np.pad(a, ((0, 0), (pad, pad)), mode="edge")94 win = np.lib.stride_tricks.sliding_window_view(ap, pool_kernel, axis=1)95 a = win.max(axis=-1) if pooling == "max" else win.mean(axis=-1)96 scores.append(a)97 return scores9899100def topk_indices_per_head(score: np.ndarray, keep: int, always_keep: np.ndarray | None = None) -> np.ndarray:101 """score [H, T] -> indices [H, keep] sorted ascending per head; always_keep positions forced in."""102 H, T = score.shape103 keep = max(min(keep, T), 1)104 out = np.zeros((H, keep), dtype=np.int64)105 for h in range(H):106 s = score[h].copy()107 if always_keep is not None and len(always_keep):108 s[always_keep] = np.inf109 idx = np.argpartition(-s, keep - 1)[:keep]110 out[h] = np.sort(idx)111 return out112113114def pyramid_budgets(total_keep: int, n_layers: int, window: int, beta: int = 20) -> list[int]:115 """PyramidKV budget rule (KVCache-Factory): per-layer budgets decreasing linearly, summing ~ total."""116 cap = total_keep117 base = max(cap - window, 1)118 mn = max(base // beta, 1)119 mxb = 2 * base - mn120 steps = (mxb - mn) // max(n_layers - 1, 1)121 budgets = [max(mxb - l * steps, mn) + window for l in range(n_layers)]122 # rescale to preserve total budget exactly123 scale = total_keep * n_layers / sum(budgets)124 budgets = [max(1, int(round(b * scale))) for b in budgets]125 return budgets126127128def select_cache(model, cache: list, policy: str, keep_frac: float, window_token_ids: list[int] | None = None,129 seed: int = 0, sinks: int = 4, pool_kernel: int = 7, pooling: str = "max", window_keep: int = 32) -> tuple[list, dict]:130 """Build a sparse cache under `policy` retaining round(keep_frac*T) positions. Returns (cache, info)."""131 T = cache[0].offset132 keep = max(1, int(round(keep_frac * T)))133 L = len(cache)134 info = {"policy": policy, "keep_frac": keep_frac, "T": T, "keep": keep}135 import time136 t0 = time.perf_counter()137 if policy == "sink_recent":138 idx = sink_recent_indices(T, keep, sinks)139 out = select_positions(cache, [idx] * L)140 elif policy == "random":141 idx = random_indices(T, keep, seed)142 out = select_positions(cache, [idx] * L)143 elif policy in ("attention", "pyramid"):144 if not window_token_ids:145 raise ValueError("attention policies need window_token_ids")146 scores = attention_window_scores(model, cache, window_token_ids, pool_kernel, pooling)147 always = np.arange(max(0, T - window_keep), T)148 if policy == "attention":149 per_layer = [topk_indices_per_head(s, keep, always) for s in scores]150 else:151 budgets = pyramid_budgets(keep, L, window_keep)152 per_layer = [topk_indices_per_head(s, b, always) for s, b in zip(scores, budgets)]153 info["layer_budgets"] = budgets154 out = select_positions(cache, per_layer)155 info["window_tokens"] = len(window_token_ids)156 else:157 raise ValueError(policy)158 mx.eval([x for c in out for x in (c.keys, c.values)])159 info["compress_ms"] = (time.perf_counter() - t0) * 1000.0160 return out, info161