# SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget

arXiv: https://arxiv.org/abs/2404.04793 · HTML: https://mlsanity.com/paper/2404.04793/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 15 of its 18 references that are in the map.
- First posted on arXiv 7 Apr 2024.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 15 of 15)

- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://mlsanity.com/paper/2306.14048/) (2023)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time](https://mlsanity.com/paper/2305.17118/) (2023)
- [Adaptive Attention Span in Transformers](https://mlsanity.com/paper/1905.07799/) (2019)
- [SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference](https://mlsanity.com/paper/2307.02628/) (2023)
- [LLMCarbon: Modeling the end-to-end Carbon Footprint of Large Language Models](https://mlsanity.com/paper/2309.14393/) (2023)
- [Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers](https://mlsanity.com/paper/2305.15805/) (2023)
- [Do Transformers Need Deep Long-Range Memory](https://mlsanity.com/paper/2007.03356/) (2020)

## Built on by (most cited first, 23 of 23)

- [PQCache: Product Quantization-based KVCache for Long Context LLM Inference](https://mlsanity.com/paper/2407.12820/) (2024)
- [Cognitive Memory in Large Language Models](https://mlsanity.com/paper/2504.02441/) (2025)
- [Dialogue Without Limits: Constant-Sized KV Caches for Extended Responses in LLMs](https://mlsanity.com/paper/2503.00979/) (2025)
- [LightTransfer: Your Long-Context LLM is Secretly a Hybrid Model with Effortless Adaptation](https://mlsanity.com/paper/2410.13846/) (2024)
- [KV Cache Compression for Inference Efficiency in LLMs: A Review](https://mlsanity.com/paper/2508.06297/) (2025)
- [ZSMerge: Zero-Shot KV Cache Compression for Memory-Efficient Long-Context LLMs](https://mlsanity.com/paper/2503.10714/) (2025)
- [Rethinking Key-Value Cache Compression Techniques for Large Language Model Serving](https://mlsanity.com/paper/2503.24000/) (2025)
- [DualMap: Enabling Both Cache Affinity and Load Balancing for Distributed LLM Serving](https://mlsanity.com/paper/2602.06502/) (2026)
- [SemShareKV: Efficient KVCache Sharing for Semantically Similar Prompts via Token-Level LSH Matching](https://mlsanity.com/paper/2509.24832/) (2025)
- [Enhancing Instruction Following of LLMs via Activation Steering with Dynamic Rejection](https://mlsanity.com/paper/2603.06745/) (2026)
- [HCAttention: Extreme KV Cache Compression via Heterogeneous Attention Computing for LLMs](https://mlsanity.com/paper/2507.19823/) (2025)
- [Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference](https://mlsanity.com/paper/2601.07667/) (2026)
- [ResKV: Reconstructing Omitted Attention Contributions for Fixed-Budget KV Cache Compression](https://mlsanity.com/paper/2607.29591/) (2026)
- [KV-Kaizen: Learning Context-Adaptive Cache Compression Choices](https://mlsanity.com/paper/2609.37988/) (2026)
- [Batch-Max: Higher LLM Throughput using Larger Batch Sizes and KV Cache Compression](https://mlsanity.com/paper/2412.05693/) (2024)
- [BaKlaVa -- Budgeted Allocation of KV cache for Long-context Inference](https://mlsanity.com/paper/2502.13176/) (2025)
- [Whole-Piece Training for Symbolic Music Language Models via Full-Horizon Compressed Recurrence](https://mlsanity.com/paper/2602.19816/) (2026)
- [WindowQuant: Mixed-Precision KV Cache Quantization based on Window-Level Similarity for VLMs Inference Optimization](https://mlsanity.com/paper/2605.02262/) (2026)
- [From Tensor Buffer to Distributed Memory Hierarchy: A Survey of KV Cache Management for LLM Serving](https://mlsanity.com/paper/2607.02574/) (2026)
- [HeadWiseKV: Budgeted Per-Head Cache Residency for Hybrid Long-Context Language Models](https://mlsanity.com/paper/2609.02029/) (2026)
- [Periodic Weak Spots: Phase Sensitivity from Chunked KV-Cache Compression](https://mlsanity.com/paper/2609.36322/) (2026)
- [JevSpawn: Adaptive Agentic Inference through Compositional Action Spaces](https://mlsanity.com/paper/2610.00437/) (2026)
- [MultiAlloc: Budget-Aware Precision Allocation for KV Cache Quantization](https://mlsanity.com/paper/2610.18642/) (2026)

## Abstract (arXiv)

Optimizing the Key-Value (KV) cache of the Large Language Model (LLM) has been considered critical to saving the cost of inference. Most of the existing KV-cache compression algorithms attempted to sparsify the sequence of tokens by taking advantage of the different importance of tokens. However, most of these methods treat all layers equally, allocating the same KV budget to each layer. This approach is suboptimal, as some layers may be less sensitive to input tokens yet still receive the same budget as others. In this work, we found that by identifying the importance of attention layers, we could optimize the KV-cache jointly from two dimensions, i.e., sequence-wise and layer-wise. Based on our observations regarding layer-wise importance in inference, we propose SqueezeAttention to precisely optimize the allocation of KV-cache budget among layers on-the-fly and then incorporate three representative sequence-wise algorithms to compress the KV-cache for each layer with its very own budget. Specifically, we first measure each layer's importance by calculating the cosine similarity of the input prompt differences before and after the self-attention layers. Based on this similarity, we then categorize the layers into two groups and adjust their KV budgets accordingly. By optimizing the KV-cache from both sequence's and layer's dimensions, SqueezeAttention achieves around 30% to 70% of the memory reductions and up to 2.2 times of throughput improvements in a wide range of LLMs and benchmarks. The code is available at https://github.com/hetailang/SqueezeAttention.