skhavin/proactive-cache

1"""2press.py — KVPress-compatible wrapper for ProactiveCache eviction.3 4Implements the BasePress API from NVIDIA's kvpress library so that5ProactiveCache can be benchmarked directly against the 20+ methods6in the KVPress standard evaluation suite.7 8Usage (requires: pip install kvpress):9    from proactive_cache import ProactiveCachePress10    press = ProactiveCachePress(compression_ratio=0.75, prototype_path="...")11    # Use with kvpress evaluation harness12"""13 14from __future__ import annotations15import os16import pickle17import torch18import numpy as np19from dataclasses import dataclass20from typing import Optional21 22from .eviction import score_tokens23 24 25# Shim for Python 3.13+ which removed the 'pipes' module (needed by fire/kvpress)26try:27    import pipes28except ImportError:29    import sys, shlex30    sys.modules['pipes'] = shlex31 32# ── KVPress ScorerPress compatibility shim ────────────────────────────────────33try:34    from kvpress import ScorerPress35    _KVPRESS_AVAILABLE = True36except ImportError:37    _KVPRESS_AVAILABLE = False38 39    class ScorerPress:40        """Minimal shim — allows import without kvpress installed."""41        def __init__(self):42            self.compression_ratio = 0.043 44        def score(self, module, hidden_states, keys, values, attentions, kwargs):45            raise NotImplementedError("Install kvpress: pip install kvpress")46 47 48@dataclass49class ProactiveCachePress(ScorerPress):50    """51    KVPress-compatible Proactive KV Cache eviction plugin.52 53    Implements the BasePress.score() hook, called once per attention layer54    during prefill. Returns a scalar importance score per token position —55    higher score = keep, lower score = evict (following KVPress convention).56 57    Args:58        compression_ratio: Fraction of tokens to EVICT [0.0, 1.0).59            e.g. 0.75 → keep 25% of the KV cache (budget = seq_len * 0.25).60        prototype_path: Path to a prototypes .pkl file from ``ProactiveCache.profile()``.61            If None, falls back to attention-sink + recency-only scoring.62 63    Example:64        press = ProactiveCachePress(compression_ratio=0.75, prototype_path="protos.pkl")65    """66    compression_ratio: float = 0.567    prototype_path: Optional[str] = None68 69    def __post_init__(self):70        self._prototypes = None71        if self.prototype_path and os.path.exists(self.prototype_path):72            with open(self.prototype_path, "rb") as f:73                self._prototypes = pickle.load(f)74            print(f"[ProactiveCachePress] Loaded {len(self._prototypes)} prototypes "75                  f"from {self.prototype_path}")76        else:77            print("[ProactiveCachePress] No prototypes loaded — using sink+recency scoring.")78 79    def score(self, module, hidden_states, keys, values, attentions, kwargs):80        """81        KVPress hook: called once per attention layer during the prefill pass.82 83        Returns:84            scores: (batch, num_heads, seq_len) float tensor.85                    Higher = more important. KVPress will keep the top-K tokens86                    where K = seq_len * (1 - compression_ratio).87        """88        batch_size, num_heads, seq_len, head_dim = keys.shape89        budget = max(1, int(seq_len * (1.0 - self.compression_ratio)))90        device = keys.device91 92        # Build position scores (O(n), query-free)93        proto_scores = score_tokens(self._prototypes, seq_len, budget)94        proto_tensor = torch.tensor(proto_scores, dtype=torch.float32, device=device)95 96        # Broadcast (1, 1, seq_len) → (batch, num_heads, seq_len)97        scores = proto_tensor.unsqueeze(0).unsqueeze(0).expand(batch_size, num_heads, seq_len)98        return scores.contiguous()99