skhavin/proactive-cache

1"""2eviction.py — Token scoring and KV cache pruning.3 4Core O(n) eviction policy:5  1. Score each token position using offline-profiled prototype centroids.6  2. Keep the top-budget tokens (attention sink + recency anchors + semantic prototypes).7  3. Prune the KV cache to exactly `budget` positions.8 9This is coordinate-free and RoPE-compatible: we only select positions, never10reorder them, so relative position encodings remain valid.11"""12 13from __future__ import annotations14import torch15import numpy as np16from typing import Optional, Dict, Tuple, List17 18from .utils import to_tuple_kv, to_dynamic_cache19 20 21def score_tokens(22    prototypes: Optional[Dict],23    seq_len: int,24    budget: int,25) -> np.ndarray:26    """27    Score all token positions using prototype centroid histograms.28 29    Algorithm (O(n) per call):30      - For each profiled (layer, head), accumulate the centroid attention31        histogram as a distance-weighted score over token positions.32      - Boost attention sink (token 0) unconditionally.33      - Boost a proportional recency window at the tail.34      - Add a small deterministic tiebreaker (position index).35 36    Args:37        prototypes: Output of ``build_prototypes()``. If None, falls back to38                    uniform scoring (no-op — keep all tokens equally).39        seq_len:    Current sequence length to score.40        budget:     Target number of tokens to keep.41 42    Returns:43        scores: (seq_len,) float64 array. Higher = more important.44    """45    scores = np.zeros(seq_len, dtype=np.float64)46 47    if prototypes is not None:48        for (layer, head), data in prototypes.items():49            centroid = data["centroids"][0]            # shape: (profile_seq_len,)50            max_d = min(len(centroid), seq_len)51            if max_d == 0:52                continue53            cumsum = np.cumsum(centroid[:max_d])54            for p in range(seq_len):55                reach = min(max_d, seq_len - p)56                if reach > 0:57                    scores[p] += cumsum[reach - 1]58 59    # ── Robust Split-Budget Boosting (Sinks + 50% Recency + 50% Semantic) ─────60    # Ensures perfect stability on relative position models (like LLaMA/RoPE)61    # by guaranteeing a large contiguous local context window and a secure sink.62    peak = scores.max() if scores.max() > 0 else 1.063 64    # 1. Boost Attention Sinks (first 4 tokens) securely65    for i in range(min(4, seq_len)):66        scores[i] += peak * 100.067 68    # 2. Boost Recency Window (50% of the budget) securely69    recency_window = min(max(8, budget // 2), seq_len)70    for i in range(recency_window):71        scores[seq_len - 1 - i] += peak * 50.072 73    # ── Deterministic tiebreaker (prefer later tokens among equals) ───────────74    scores += np.linspace(0, 1e-4, seq_len)75 76    return scores77 78 79def select_indices(scores: np.ndarray, budget: int) -> List[int]:80    """Return the top-budget indices, sorted in ascending order (preserves sequence order)."""81    actual_budget = min(budget, len(scores))82    top = np.argsort(scores)[-actual_budget:]83    return sorted(top.tolist())84 85 86def prune_kv_cache(87    past_key_values,88    indices: List[int],89    device: torch.device,90):91    """92    Prune a KV cache to the given token indices.93 94    Args:95        past_key_values: DynamicCache or legacy tuple from a model forward pass.96        indices:         Sorted list of token indices to keep.97        device:          CUDA/CPU device for the index tensor.98 99    Returns:100        Pruned KV cache in the same format the model expects101        (DynamicCache if transformers ≥ 4.38, else tuple).102    """103    idx_t = torch.tensor(indices, dtype=torch.long, device=device)104    kv_tuple = to_tuple_kv(past_key_values)105    pruned = tuple(106        (k.index_select(2, idx_t), v.index_select(2, idx_t))107        for k, v in kv_tuple108    )109    return to_dynamic_cache(pruned)110 111 112def evict(113    past_key_values,114    budget: int,115    prototypes: Optional[Dict],116    seq_len: int,117    device: torch.device,118):119    """120    One-shot eviction: score → select → prune.121 122    If ``seq_len <= budget``, returns ``past_key_values`` unchanged.123    """124    if seq_len <= budget:125        return past_key_values126 127    scores = score_tokens(prototypes, seq_len, budget)128    indices = select_indices(scores, budget)129    return prune_kv_cache(past_key_values, indices, device)130