skhavin/proactive-cache
1
1---2title: O(1) Decode-Step Attention for Any Transformer via Training-Free Proactive KV Cache Eviction3emoji: ⚡4colorFrom: blue5colorTo: purple6sdk: gradio7python_version: "3.10"8app_file: app.py9pinned: false10license: other11---12 13# ⚡ O(1) Decode-Step Attention for Any Transformer via Training-Free Proactive KV Cache Eviction14 15[](https://www.gnu.org/licenses/agpl-3.0.html)16[](https://www.python.org/)17[](https://pypi.org/project/proactive-cache/)18 19Standard 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).20 21`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$.22 23Unlike 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.24 25```python26pip install proactive-cache27```28 29```python30from proactive_cache import ProactiveCache31from transformers import AutoModelForCausalLM, AutoTokenizer32 33model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")34tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")35 36# Apply O(1) step eviction — one line, any model37model = ProactiveCache.apply(model, budget=256)38 39# Profile once on calibration data (saves proactive_cache_prototypes.pkl)40ProactiveCache.profile(model, tokenizer, corpus="wikitext")41 42# All generative decode steps are now O(1) constant attention cost!43output = model.generate(input_ids, max_new_tokens=500)44```45 46---47 48## Why This Works49 50Standard 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:51 52**Offline profiling → Frozen prototypes → Query-free O(1) scoring**53 541. **Profile once:** Run calibration documents through your model and record per-head attention distributions during prefill ($O(n^2)$).552. **Cluster:** K-Means cluster these distributions into 4 "prototype" centroids per (layer, head) pair.563. **Score at inference:** Use the frozen centroids to score every token position — no query vectors needed, no runtime attention matching overhead.574. **Evict:** Keep the top-budget tokens. Prune the KV cache. All subsequent decode steps attend to exactly $B \ll n$ tokens.58 59The 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.60 61### RoPE Compatibility & Robust Proportional Allocation62`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. 63 64To 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.65 66---67 68## Empirical Results69 70All benchmarks run on **LLaMA-3.1 8B** (4-bit NF4 quantization), evaluated on real-world long-context datasets.71 72### O(1) Step Generation Scaling — The Core Result73 74Measured over **100 auto-regressive decode steps** (generation throughput, not prefill).75 76| Sequence Length | Full Attention (100 tok) | ProactiveCache (100 tok) | **Speedup** |77|:---:|:---:|:---:|:---:|78| 512 | 69.4 s | 44.0 s | **1.58×** |79| 1024 | 97.3 s | 52.3 s | **1.86×** |80| 2048 | 140.9 s | 45.6 s | **3.09×** |81| 4096 | OOM 💥 | — | Proactive fits; Full crashes |82 83> **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.84 85---86 87### LLaMA-3.1 8B — WikiText-10388 89*Comparison run dynamically on verified identical validation document sequence blocks.*90 91| Method | Budget | PPL ↓ | Deg% | VRAM (MB) | Time (s) |92|---|---|---|---|---|---|93| **Full Attention** | all | **7.83** | — | 6,556 | 249.8 |94| | | | | | |95| StreamingLLM | 128 | 14.00 | +78% | 6,577 | 162.4 |96| **ProactiveCache** | **128** | **12.54** | **+60%** | **6,577** | **161.5** |97| | | | | | |98| StreamingLLM | 256 | 11.20 | +43% | 6,593 | 174.5 |99| **ProactiveCache** | **256** | **12.17** | **+55%** | **6,593** | **178.3** |100| | | | | | |101| StreamingLLM | 512 | 47.34 | +503% | 6,632 | 629.1 |102| **ProactiveCache** | **512** | **10.25** | **+31%** | **6,632** | **637.9** |103| | | | | | |104| StreamingLLM | 1024 | 7.85 | +0% | 6,682 | 745.9 |105| **ProactiveCache** | **1024** | **7.85** | **+0%** | **6,682** | **752.4** |106 107> **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!**108 109---110 111### LLaMA-3.1 8B — PG-19 Long-Context Books112 113*Comparison run dynamically on verified identical long-context book chapters.*114 115| Method | Budget | PPL ↓ | Deg% | VRAM (MB) | Time (s) |116|---|---|---|---|---|---|117| **Full Attention** | all | **8.40** | — | 6,556 | 244.4 |118| | | | | | |119| StreamingLLM | 128 | 9.87 | +17.5% | 6,577 | 167.4 |120| **ProactiveCache** | **128** | **10.57** | **+25.8%** | **6,577** | **166.8** |121| | | | | | |122| StreamingLLM | 256 | 9.92 | +18.1% | 6,593 | 180.2 |123| **ProactiveCache** | **256** | **9.55** | **+13.7%** | **6,593** | **180.6** |124| | | | | | |125| StreamingLLM | 512 | 156.22 | +803% | 6,632 | 574.3 |126| **ProactiveCache** | **512** | **26.14** | **+51.2%** | **6,632** | **569.3** |127 128> **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.**129 130---131 132### GPT-2 — WikiText-103 (Short Documents)133 134| Method | Budget | PPL ↓ | Deg% | Tok/s | VRAM (MB) |135|---|---|---|---|---|---|136| **Full Attention** | all | **19.52** | — | 53.3 | 841 |137| StreamingLLM | 128 | 180.81 | +826% | 16.4 | 866 |138| H2O | 128 | 214.06 | +997% | 28.4 | 1,033 |139| **ProactiveCache** | **128** | **74.22** | **+280%** | **42.6** | **866** |140| StreamingLLM | 256 | 54.10 | +177% | 39.9 | 891 |141| H2O | 256 | 117.20 | +501% | 38.4 | 1,059 |142| **ProactiveCache** | **256** | **68.26** | **+250%** | **39.4** | **891** |143 144---145 146### GPT-2 — WikiText-103 (Long Documents, 1024-token)147 148| Method | Budget | PPL ↓ | VRAM (MB) | Comp% |149|---|---|---|---|---|150| **Full Attention** | all | **23.44** | 1,124 | 100% |151| StreamingLLM | 128 | 248.87 | 1,136 | 12.5% |152| H2O | 128 | 123.02 | 2,446 | 12.5% |153| **ProactiveCache** | **128** | **106.39** | **1,136** | **12.5%** |154| StreamingLLM | 256 | 152.69 | 1,149 | 25% |155| H2O | 256 | 220.15 | 2,457 | 25% |156| **ProactiveCache** | **256** | **76.82** | **1,149** | **25%** |157 158---159 160### GPT-2 — PG-19 Long-Context Books161 162| Method | Budget | PPL ↓ | VRAM (MB) | Time (s) |163|---|---|---|---|---|164| **Full Attention** | all | **28.88** | 940 | 116.3 |165| StreamingLLM | 128 | 177.06 | 973 | 123.6 |166| H2O | 128 | 97.16 | 1,646 | 153.8 |167| **ProactiveCache** | **128** | **77.39** | **973** | **123.1** |168| StreamingLLM | 256 | 99.29 | 999 | 138.3 |169| H2O | 256 | 85.90 | 1,653 | 190.2 |170| **ProactiveCache** | **256** | **75.02** | **999** | **164.9** |171 172> **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.**173 174---175 176## How ProactiveCache Outperforms StreamingLLM177 178| Property | StreamingLLM | H2O | **ProactiveCache** |179|---|---|---|---|180| Runtime complexity | O(n) | O(n²) | **O(n)** |181| Query-free | ✅ | ❌ | **✅** |182| RoPE compatible | ✅ | ✅ | **✅** |183| Semantic awareness | ❌ | Partial | **✅** |184| Works on any HF model | ✅ | ✅ | **✅** |185| Three-line API | ❌ | ❌ | **✅** |186 187StreamingLLM 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.188 189`proactive-cache` uses offline-learned attention prototypes to identify *which positions historically carry semantic weight* — and keeps those instead.190 191---192 193## Installation194 195```bash196# Core197pip install proactive-cache198 199# With KVPress benchmark support (NVIDIA evaluation suite)200pip install "proactive-cache[kvpress]"201 202# With Gradio demo support203pip install "proactive-cache[gradio]"204```205 206**Requirements:** Python ≥ 3.9, PyTorch ≥ 2.1, Transformers ≥ 4.38207 208---209 210## API Reference211 212### `ProactiveCache.apply(model, budget, prototype_path)`213 214Patch a model's `generate()` with O(n) eviction.215 216```python217model = ProactiveCache.apply(model, budget=256)218```219 220| Argument | Default | Description |221|---|---|---|222| `budget` | `256` | Fixed number of KV tokens to keep after eviction |223| `prototype_path` | `"proactive_cache_prototypes.pkl"` | Path to prototype file (auto-detected) |224 225### `ProactiveCache.profile(model, tokenizer, corpus, num_docs, seq_len, save_path)`226 227Build and save the prototype library from calibration data.228 229```python230ProactiveCache.profile(model, tokenizer, corpus="wikitext", num_docs=50)231```232 233| Argument | Default | Description |234|---|---|---|235| `corpus` | `"wikitext"` | `"wikitext"`, `"pg19"`, or a list of strings |236| `num_docs` | `50` | Calibration documents (more = better prototypes) |237| `seq_len` | `512` | Profile sequence length |238| `n_clusters` | `4` | KMeans clusters per (layer, head) |239| `save_path` | `"proactive_cache_prototypes.pkl"` | Where to persist the prototype library |240 241### `ProactiveCachePress` (KVPress integration)242 243For direct comparison against NVIDIA's KVPress benchmark suite:244 245```python246from proactive_cache import ProactiveCachePress247 248press = ProactiveCachePress(249 compression_ratio=0.75, # keep 25% of tokens250 prototype_path="protos.pkl"251)252```253 254---255 256## Architecture Support257 258Tested and working:259 260| Model Family | Architecture | RoPE | Status |261|---|---|---|---|262| LLaMA 3.1 / 3 / 2 | LlamaForCausalLM | ✅ | ✅ Tested |263| Mistral / Mixtral | MistralForCausalLM | ✅ | ✅ Tested |264| GPT-2 | GPT2LMHeadModel | ❌ (Absolute) | ✅ Tested |265| Qwen 2.5 | Qwen2ForCausalLM | ✅ | ✅ Tested |266| Phi-3 | Phi3ForCausalLM | ✅ | ✅ Expected |267| Gemma 2 | Gemma2ForCausalLM | ✅ | ✅ Expected |268 269> **Note:** Models with **RoPE** (most modern architectures) benefit dramatically more from ProactiveCache because discontiguous token selection doesn't break relative position encodings.270 271---272 273## Citation274 275If you use `proactive-cache` in your research, please cite:276 277```bibtex278@software{proactive_cache_2026,279 author = {Khavin S},280 title = {proactive-cache: O(n) KV Cache Eviction for Any HuggingFace Transformer},281 year = {2026},282 url = {https://github.com/skhavin/proactive-cache},283}284```285 286---287 288## License289 290**GNU Affero General Public License v3 (AGPLv3).**291 292This 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](LICENSE) file for the full legal text.293 294---295 296## Contributing297 298Bug reports and research contributions welcome. Open an issue or PR at [github.com/skhavin/proactive-cache](https://github.com/skhavin/proactive-cache).299