Metadata-Version: 2.4
Name: hippocache
Version: 0.1.2
License-File: LICENSE
Summary: Smart KV cache for transformer inference
Requires-Python: >=3.8
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Homepage, https://github.com/AdvaitSan/HippoCache

# hippocache 🦛

Smart, plug-and-play KV cache compression for Transformer inference, written in Rust.

`hippocache` provides drop-in KV cache implementations with multiple retention and compression strategies, designed to fit into LLM generation loops to reduce memory footprints without sacrificing model perplexity.

---

## 🚀 Features

- 🦀 **Rust Engine**: High-performance token merging and eviction operations implemented using `ndarray` and PyO3.
- 🔌 **Drop-in Integration**: Fits seamlessly into standard LLM autoregressive generation loops.
- 📊 **Smart Compression**: Includes advanced strategies like cosine-similarity-based sequence merging (`merge` and `hybrid`).
- 📈 **Stats & Metrics**: Built-in tracking for compression ratios, merges, and evictions.

---

## 📦 Installation

Install the pre-compiled package directly via `pip`:

```bash
pip install hippocache
```

---

## ⚡ Quickstart

Using `hippocache` is simple. Below is a realistic example demonstrating how to integrate `KVCache` into a token generation loop:

```python
import random
from hippocache import KVCache

# 1. Initialize the KV Cache
# d_k and d_v represent the key and value head dimensions
cache = KVCache(
    d_k=64,
    d_v=64,
    strategy="hybrid",
    static_k=3,         # Retain the first 3 tokens (e.g., prompt context)
    threshold=0.90,     # Cosine similarity merge threshold
    merge_every=4       # Check and merge adjacent tokens every 4 steps
)

# 2. Simulate a generation loop
for step in range(20):
    # Mock key & value vectors for the new token
    new_k = [random.uniform(-1.0, 1.0) for _ in range(64)]
    new_v = [random.uniform(-1.0, 1.0) for _ in range(64)]
    
    # Append to cache (merges/evictions happen automatically in Rust)
    cache.append(new_k, new_v)
    
    # Query attention over the compressed cache
    q = [random.uniform(-1.0, 1.0) for _ in range(64)]
    attention_output = cache.attend(q)
    
    print(f"Step {step+1:02d} | Cache Size: {cache.current_len()} | Stats: {cache.stats()}")

# 3. Retrieve final stats
final_stats = cache.stats()
print("\nFinal Generation Stats:")
print(f"  Strategy:          {final_stats['strategy']}")
print(f"  Compression Ratio: {final_stats['compression_ratio']:.2f} (lower is better)")
print(f"  Total Appends:     {final_stats['total_appends']}")
print(f"  Total Merges:      {final_stats['total_merges']}")
```

---

## 🛠 Caching Strategies

| Strategy | Config Parameters | Description |
| :--- | :--- | :--- |
| `full` | None | Keeps all key-value states in memory (default). |
| `sliding` | `window_size` | Retains only the last $N$ tokens, evicting older states. |
| `static` | `static_k`, `recent_k` | Keeps the first $K$ prefix tokens forever + the last $R$ recent tokens. |
| `merge` | `threshold`, `merge_every` | Compares adjacent key tensors using cosine similarity and merges them if similarity exceeds the threshold. |
| `hybrid` | `static_k`, `threshold`, `merge_every` | Keeps the first $K$ prefix tokens forever and performs similarity merges on the rest of the sequence. |

---

## 📊 Benchmark Results

Evaluated on **TinyLlama-1.1B** (Layer 0, Head 0) across various compression strategies and parameters.

<p align="center">
  <img src="https://hippo-cache.vercel.app/assets/overview-tLeBr0jL.png" alt="HippoCache Compression Analysis" width="100%" />
</p>

### 1. Strategy Comparison
*Evaluating the tradeoff between cache size and quality loss (measured by KL Divergence / proxy perplexity and L2 attention divergence).*

<p align="center">
  <img src="https://hippo-cache.vercel.app/assets/strat_cmrps-B03pWHuf.png" alt="Strategy Comparison - Cache Size vs Quality" width="100%" />
</p>

| Strategy | Config | Cache Size (Tokens) | Size Reduction (%) | L2 Attention Divergence | KL Divergence (Quality Loss) |
| :--- | :--- | :---: | :---: | :---: | :---: |
| **`full`** | (Default) | 16 | 0% | 0.0000 | 0.0000 |
| **`merge`** | `threshold=0.90` | 12 | **25%** | **0.0031** | **0.67 × 10⁻⁵** |
| **`hybrid`** | `static_k=3, threshold=0.90` | 12 | **25%** | **0.0032** | **0.74 × 10⁻⁵** |
| **`static`** | `static_k=3, recent_k=6` | 9 | 43.8% | 0.0051 | 1.95 × 10⁻⁵ |
| **`sliding`** | `window_size=8` | 8 | 50% | 0.0076 | 3.75 × 10⁻⁵ |

> [!TIP]
> **Key Insight**: The `merge` and `hybrid` strategies achieve **over 25% cache size reduction** with **3x to 5x less quality loss** (lower KL divergence) compared to traditional static or sliding window strategies at similar sizes.

---

### 2. Compression vs. Quality (Merge Threshold Sweep)
*Adjusting the cosine similarity threshold allows fine-grained control over the compression vs. attention quality tradeoff.*

<p align="center">
  <img src="https://hippo-cache.vercel.app/assets/threshold_sweep-CFv4A1yX.png" alt="Threshold Sweep - Variance Across Heads" width="100%" />
</p>

| Merge Threshold | Cache Size (Tokens) | Cache Compression | L2 Attention Divergence | Recommendation |
| :---: | :---: | :---: | :---: | :--- |
| **`0.99`** | 16 / 16 | 0% | 0.0000 | Lossless baseline |
| **`0.97`** | 15 / 16 | 6.2% | 0.0018 | Ultra-conservative |
| **`0.95`** | 15 / 16 | 6.2% | 0.0018 | Ultra-conservative |
| **`0.92`** | 14 / 16 | 12.5% | 0.0032 | Conservative |
| **`0.90`** | 12 / 16 | **25.0%** | **0.0030** | **Sweet spot (recommended)** |
| **`0.85`** | 8 / 16 | 50.0% | 0.0051 | Moderate compression |
| **`0.80`** | 4 / 16 | 75.0% | 0.0140 | High compression, quality degradation |

---

## 📜 License

MIT License. See [LICENSE](file:///data/RUST/hippocache/LICENSE) for details.

