Research map / Architectures and efficiency / Attention and transformers
SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SqueezeAttention, a 2D KV-cache compression method that reallocates per-layer cache budgets using cosine-similarity-based layer importance clustering, combined with existing sequence-wise eviction policies.
Problem
KV-cache compression methods for LLM inference allocate the same token budget to every attention layer, which is suboptimal because layers differ in importance.
Approach
Measure each layer's importance during prefill by the cosine similarity between hidden states before and after self-attention, cluster layers into three groups with KMeans, and shift budget from less important layers to important ones while keeping the total budget fixed; any sequence-wise eviction policy can then be applied per layer.
Key results
- SqueezeAttention reaches comparable accuracy with 20% KV budget for Mistral-7B on SAMSUM (41.05 vs 43.53 full cache) versus 30% for the best baseline.
- Up to 2.2x throughput improvement for Mistral-7B over full cache at batch size 128 (824.4 vs OOM-limited full cache), and 1.4x for Llama2-70B.
- Prefilling overhead is 6.3% for Mistral-7B at 8k prompt length on one A100.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add layer-wise KV-budget reallocation on top of Sliding Window | Mistral-7B · Mistral-7B, batch size 32, 512 + 1024 | throughput (token / s): 404.5 -> 496.5 | Helped |
| Add layer-wise KV-budget reallocation on top of Sliding Window | Mistral-7B · Mistral-7B, batch size 64, 512 + 1024 | throughput (token / s): 512.2 -> 682.7 | Helped |
| Add layer-wise KV-budget reallocation on top of StreamingLLM | Llama2-7B · Llama2-7B, batch size 32, 512 + 1024 | throughput (token / s): 113.7 -> 143.0 | Helped |
| Add layer-wise KV-budget reallocation on top of StreamingLLM | Llama2-70B · Llama2-70B, batch size 64, 256 + 512 | throughput (token / s): OOM -> 170.7 | Helped |
| Add layer-wise KV-budget reallocation (ablation) | SqueezeAttention · Mistral-7B, SAMSUM | ROUGE-L: 15.7; 34.01; 37.26; 37.69; 27.41; 26.29; 11.11; 10.48; 8.72; 9.07 | Mixed |
| Set p to 0.4 rather than p=1.0 (ablation) | SqueezeAttention on Mistral-7B · SAMSUM | ROUGE-L: 9.07 -> 37.69 | Helped |
| Set p to 0.1 rather than p=0.4 (ablation) | SqueezeAttention on Mistral-7B · SAMSUM | ROUGE-L: 37.69 -> 15.7 | Hurt |
Surprising findings
- Reallocating budget can increase accuracy over the fixed-budget baseline at the same total budget, e.g., Mistral-7B ROUGE-L peaks near p=0.4 and drops sharply below p=0.2 or above p=0.6.
Limitations
- Assumes a sequence-wise eviction policy exists that does not hurt accuracy under tolerance for the given model and dataset.
- Layer importance varies across tasks, so the grouping is task-dependent and must be recomputed on the fly.
- Hyperparameter p (0.3-0.4 suggested) requires tuning per setting.
Experimental setup
Datasets: CNN/Daily Mail XSUM SAMSUM NarrativeQA TriviaQA LongBench
Benchmarks: CNN/Daily Mail XSUM SAMSUM NarrativeQA TriviaQA LongBench
Metrics: Rouge-2 Rouge-L F1 memory usage throughput (token/s) prefilling time
Baselines: Heavy-Hitter Oracle (H2O) Sliding Window Attention StreamingLLM Full Cache
Builds on: H2O StreamingLLM FastGen Early-exiting LLM
Closest papers by meaning
- MiniCache: KV Cache Compression in Depth Dimension for Large Language Models (NeurIPS 2024 · similarity 0.94)
- Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference (NeurIPS 2025 · similarity 0.94)
- Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference (ICML 2024 · similarity 0.93)
- ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference (NeurIPS 2025 · similarity 0.93)
- RazorAttention: Efficient KV Cache Compression Through Retrieval Heads (ICLR 2025 · similarity 0.93)
- KVTuner: Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache Quantization for Efficient and Nearly Lossless LLM Inference (ICML 2025 · similarity 0.93)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.