# Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

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

## Key facts

- Cited by 126 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 46 of its 66 references that are in the map.
- First posted on arXiv 14 Mar 2024.

## Placed in

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

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [BERT4Rec: Sequential Recommendation with Bidirectional Encoder Representations from Transformer](https://mlsanity.com/paper/1904.06690/) (2019)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [Efficient Attentions for Long Document Summarization](https://mlsanity.com/paper/2104.02112/) (2021)

## Built on by (most cited first, 30 of 126)

- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving](https://mlsanity.com/paper/2407.00079/) (2024)
- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [LoongServe: Efficiently Serving Long-Context Large Language Models with Elastic Sequence Parallelism](https://mlsanity.com/paper/2404.09526/) (2024)
- [MagicPIG: LSH Sampling for Efficient LLM Generation](https://mlsanity.com/paper/2410.16179/) (2024)
- [Keep the Cost Down: A Review on Methods to Optimize LLM' s KV-Cache Consumption](https://mlsanity.com/paper/2407.18003/) (2024)
- [PQCache: Product Quantization-based KVCache for Long Context LLM Inference](https://mlsanity.com/paper/2407.12820/) (2024)
- [Eigen Attention: Attention in Low-Rank Space for KV Cache Compression](https://mlsanity.com/paper/2408.05646/) (2024)
- [LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference](https://mlsanity.com/paper/2407.14057/) (2024)
- [ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression](https://mlsanity.com/paper/2412.03213/) (2024)
- [Layer-Condensed KV Cache for Efficient Inference of Large Language Models](https://mlsanity.com/paper/2405.10637/) (2024)
- [LLM Inference Serving: Survey of Recent Advances and Opportunities](https://mlsanity.com/paper/2407.12391/) (2024)
- [VL-Cache: Sparsity and Modality-Aware KV Cache Compression for Vision-Language Model Inference Acceleration](https://mlsanity.com/paper/2410.23317/) (2024)
- [KVTuner: Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache Quantization for Efficient and Nearly Lossless LLM Inference](https://mlsanity.com/paper/2502.04420/) (2025)
- [The Cost of Dynamic Reasoning: Demystifying AI Agents and Test-Time Scaling from an AI Infrastructure Perspective](https://mlsanity.com/paper/2506.04301/) (2025)
- [How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad](https://mlsanity.com/paper/2406.06467/) (2024)
- [Cache-Craft: Managing Chunk-Caches for Efficient Retrieval-Augmented Generation](https://mlsanity.com/paper/2502.15734/) (2025)
- [DAPE: Data-Adaptive Positional Encoding for Length Extrapolation](https://mlsanity.com/paper/2405.14722/) (2024)
- [Tensor Product Attention Is All You Need](https://mlsanity.com/paper/2501.06425/) (2025)
- [SeerAttention-R: Sparse Attention Adaptation for Long Reasoning](https://mlsanity.com/paper/2506.08889/) (2025)
- [Cognitive Memory in Large Language Models](https://mlsanity.com/paper/2504.02441/) (2025)
- [V2PE: Improving Multimodal Long-Context Capability of Vision-Language Models with Variable Visual Position Encoding](https://mlsanity.com/paper/2412.09616/) (2024)
- [DiffKV: Differentiated Memory Management for Large Language Models with Parallel KV Compaction](https://mlsanity.com/paper/2412.03131/) (2024)
- [CriticalKV: Optimizing KV Cache Eviction from an Output Perturbation Perspective](https://mlsanity.com/paper/2502.03805/) (2025)
- [Jenga: Effective Memory Management for Serving LLM with Heterogeneity](https://mlsanity.com/paper/2503.18292/) (2025)
- [Dialogue Without Limits: Constant-Sized KV Caches for Extended Responses in LLMs](https://mlsanity.com/paper/2503.00979/) (2025)
- [Semantic Integrity Matters: Benchmarking and Preserving High-Density Reasoning in KV Cache Compression](https://mlsanity.com/paper/2502.01941/) (2025)
- [M-ANT: Efficient Low-bit Group Quantization for LLMs via Mathematically Adaptive Numerical Type](https://mlsanity.com/paper/2502.18755/) (2025)
- [xKV: Cross-Layer KV-Cache Compression via Aligned Singular Vector Extraction](https://mlsanity.com/paper/2503.18893/) (2025)
- [LazyEviction: Lagged KV Eviction with Attention Pattern Observation for Efficient Long Reasoning](https://mlsanity.com/paper/2506.15969/) (2025)

## Abstract (arXiv)

Transformers have emerged as the underpinning architecture for Large Language Models (LLMs). In generative language models, the inference process involves two primary phases: prompt processing and token generation. Token generation, which constitutes the majority of the computational workload, primarily entails vector-matrix multiplications and interactions with the Key-Value (KV) Cache. This phase is constrained by memory bandwidth due to the overhead of transferring weights and KV cache values from the memory system to the computing units. This memory bottleneck becomes particularly pronounced in applications that require long-context and extensive text generation, both of which are increasingly crucial for LLMs. This paper introduces "Keyformer", an innovative inference-time approach, to mitigate the challenges associated with KV cache size and memory bandwidth utilization. Keyformer leverages the observation that approximately 90% of the attention weight in generative inference focuses on a specific subset of tokens, referred to as "key" tokens. Keyformer retains only the key tokens in the KV cache by identifying these crucial tokens using a novel score function. This approach effectively reduces both the KV cache size and memory bandwidth usage without compromising model accuracy. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT, which employ various positional embedding algorithms. Our assessment encompasses a variety of tasks, with a particular emphasis on summarization and conversation tasks involving extended contexts. Keyformer's reduction of KV cache reduces inference latency by 2.1x and improves token generation throughput by 2.4x, while preserving the model's accuracy.