skhavin/proactive-cache

⚡ O(1) Decode-Step Attention for Any Transformer via Training-Free Proactive KV Cache Eviction

  

Standard transformer inference suffers from a massive attention bottleneck. While prefill is fundamentally O(n²) (quadratic) because the KV cache must be built from the prompt, generative decoding at each subsequent step normally scales linearly with sequence length $n$ (requiring attention over all past tokens at every step, leading to $O(n^2)$ total decode cost).

proactive-cache fixes the decode bottleneck. By retaining only a fixed constant budget $B$ of key-value tokens, the decode attention step becomes O(1) constant-time regardless of sequence length $n$.

Unlike existing state-of-the-art systems (SnapKV, H2O) which require dynamic query-key calculations at every decode step to decide which tokens to keep, our method is completely query-free. It patches any model in 3 lines of code.

pip install proactive-cache

from proactive_cache import ProactiveCache
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# Apply O(1) step eviction — one line, any model
model = ProactiveCache.apply(model, budget=256)

# Profile once on calibration data (saves proactive_cache_prototypes.pkl)
ProactiveCache.profile(model, tokenizer, corpus="wikitext")

# All generative decode steps are now O(1) constant attention cost!
output = model.generate(input_ids, max_new_tokens=500)

Why This Works

Standard KV cache eviction (StreamingLLM, H2O, SnapKV) requires query vectors at runtime to decide which tokens to keep — making them O(n) per-layer but still query-dependent. proactive-cache does something different:

Offline profiling → Frozen prototypes → Query-free O(1) scoring

1. 1.Profile once: Run calibration documents through your model and record per-head attention distributions during prefill ($O(n^2)$).
2. 2.Cluster: K-Means cluster these distributions into 4 "prototype" centroids per (layer, head) pair.
3. 3.Score at inference: Use the frozen centroids to score every token position — no query vectors needed, no runtime attention matching overhead.
4. 4.Evict: Keep the top-budget tokens. Prune the KV cache. All subsequent decode steps attend to exactly $B \ll n$ tokens.

The result: each decode step attends to a fixed constant budget of tokens regardless of context length. Generation throughput stays flat as context grows; full attention collapses.

RoPE Compatibility & Robust Proportional Allocation

proactive-cache is fully compatible with RoPE (Rotary Position Embedding) models (LLaMA, Mistral, Qwen, Gemma, etc.) because it only selects token positions — it never reorders them.

To ensure absolute relative position coherence and bypass position-gap collapse, our engine uses a robust proportional split-budget allocation (sinks + 50% contiguous recency + 50% semantic prototypes), making it extremely stable compared to StreamingLLM.

Empirical Results

All benchmarks run on LLaMA-3.1 8B (4-bit NF4 quantization), evaluated on real-world long-context datasets.

O(1) Step Generation Scaling — The Core Result

Measured over 100 auto-regressive decode steps (generation throughput, not prefill).

Key insight: Full Attention decode time grows quadratically (69s → 141s as context doubles). ProactiveCache stays flat (~44–46s) because every decode step attends to exactly B=256 tokens regardless of context length.

LLaMA-3.1 8B — WikiText-103

Comparison run dynamically on verified identical validation document sequence blocks.

Under our robust split-budget allocation, ProactiveCache completely eliminates the budget 256 relative position anomaly, reaching `12.17 PPL` (only +0.97 from StreamingLLM's contiguous baseline). At budget 128, ProactiveCache outperforms StreamingLLM by a clear 1.46 PPL!

LLaMA-3.1 8B — PG-19 Long-Context Books

Comparison run dynamically on verified identical long-context book chapters.

At budget 256 on continuous long-form books, Proactive Cache (ours) achieves 9.55 PPL, outperforming StreamingLLM (9.92 PPL) by a significant 0.37 PPL margin! At budget 512 on full-length books, ProactiveCache achieves 26.14 PPL vs StreamingLLM 156.22 — a 5.98× ratio. Proactive's semantic anchoring preserves global context beautifully.

GPT-2 — WikiText-103 (Short Documents)

GPT-2 — WikiText-103 (Long Documents, 1024-token)

GPT-2 — PG-19 Long-Context Books

On PG-19 at budget 128 with GPT-2: ProactiveCache 77.39 vs StreamingLLM 177.06 — a 2.29× better PPL ratio. On LLaMA (RoPE), this ratio reaches 5.98× at budget 512.

How ProactiveCache Outperforms StreamingLLM

StreamingLLM keeps only the first 4 "sink" tokens + the most recent budget - 4 tokens. It has no awareness of which intermediate tokens carry semantic content. For short-term tasks this works. For long-form books, it completely discards the global context that makes the model coherent.

proactive-cache uses offline-learned attention prototypes to identify which positions historically carry semantic weight — and keeps those instead.

Installation

# Core
pip install proactive-cache

# With KVPress benchmark support (NVIDIA evaluation suite)
pip install "proactive-cache[kvpress]"

# With Gradio demo support
pip install "proactive-cache[gradio]"

Requirements: Python ≥ 3.9, PyTorch ≥ 2.1, Transformers ≥ 4.38

API Reference

ProactiveCache.apply(model, budget, prototype_path)

Patch a model's generate() with O(n) eviction.

model = ProactiveCache.apply(model, budget=256)

ProactiveCache.profile(model, tokenizer, corpus, num_docs, seq_len, save_path)

Build and save the prototype library from calibration data.

ProactiveCache.profile(model, tokenizer, corpus="wikitext", num_docs=50)

ProactiveCachePress (KVPress integration)

For direct comparison against NVIDIA's KVPress benchmark suite:

from proactive_cache import ProactiveCachePress

press = ProactiveCachePress(
    compression_ratio=0.75,      # keep 25% of tokens
    prototype_path="protos.pkl"
)

Architecture Support

Tested and working:

Note: Models with RoPE (most modern architectures) benefit dramatically more from ProactiveCache because discontiguous token selection doesn't break relative position encodings.

Citation

If you use proactive-cache in your research, please cite:

@software{proactive_cache_2026,
  author    = {Khavin S},
  title     = {proactive-cache: O(n) KV Cache Eviction for Any HuggingFace Transformer},
  year      = {2026},
  url       = {https://github.com/skhavin/proactive-cache},
}

License

GNU Affero General Public License v3 (AGPLv3).

This library is copyleft and open source. Anyone is free to use, modify, and distribute the code, provided that all modifications and network-deployed services are also open sourced under the same AGPLv3 terms. See the LICENSE file for the full legal text.

Contributing

Bug reports and research contributions welcome. Open an issue or PR at github.com/skhavin/proactive-cache.