hippocache 0.1.1

Smart KV cache for transformer inference

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:

pip install hippocache

---

⚡ Quickstart

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

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.

1. Strategy Comparison

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

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.

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 for details.