Source-linked AI summary
Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, Jianfeng Gao
TL;DR
Large language model inference faces increasing KV-cache memory demands, especially as model size and generation length grow. FastGen profiles attention-head structures and adaptively compresses each head’s cache without fine-tuning, achieving substantial compression with negligible generation-quality loss. The method’s reported evaluations show strong cache reduction across tasks, while its scope includes assumptions and architectural boundaries identified by the authors.
Problem
Growing model size and generation length make KV-cache memory costly, while offloading can add inference overhead; compression without costly retraining or fine-tuning is needed.
Method
FastGen profiles attention heads and adaptively selects cache policies that preserve their observed attention structures.
Results
FastGen recovers over 95% of attention scores with 35% cache compressed and, on Llama 1 30B, surpasses fixed methods at 50% versus 15% cache compression.
Takeaways & Limitations
FastGen provides a plug-and-play approach for reducing KV-cache memory with negligible generation-quality loss and no fine-tuning.
Takeaways & Limitations
FastGen assumes attention-map structure remains stable during generation and uses the original multi-head architecture, leaving grouped-query attention integration for future work.
Abstract
from arXiv · showhide
In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.
1 INTRODUCTION
FastGen addresses the growing memory burden of KV caching by profiling attention-head structures and adapting compression policies accordingly. Across diverse generative tasks, it compresses KV caches with negligible quality loss and requires no fine-tuning.
- Motivation: KV-cache memory grows with model size and generation length, creating GPU-memory pressure and motivating compression without costly retraining or fine-tuning.Offloading can add inference overhead because of limited PCIe bandwidth.
- Approach: FastGen profiles attention heads and constructs their KV caches adaptively according to the identified structures.The approach uses a diagnose-before-compress workflow during inference preparation.
- Adaptive policies: Different policies evict long-range contexts, retain special tokens or punctuation, keep frequently attended tokens, or store all tokens for broadly attending heads.These policies correspond to local, special-token, column-sparse, and broad-attention structures.
- Deployment: FastGen is plug-and-play and does not require fine-tuning, avoiding the extra training cost associated with deploying compression on extra-large models.The method is designed as a deployment-compatible compression approach.
- Results: FastGen recovers over 95% of attention scores with 35% cache compressed and, for Llama 1 30B, surpasses fixed methods at 50% versus 15% cache compression.The evaluation covers math, code, knowledge, and common-sense reasoning tasks.
2 RELATED WORK
Prior efficiency methods compress tokens or study attention structure, but many are unsuitable for autoregressive LLMs because they target non-autoregressive models or require additional retraining. Recent KV-cache work motivates pruning cached tokens for autoregressive generation.
- Token Dropping and KV Cache Compression: Earlier token-dropping methods remove redundant words or compress sequences, but are generally designed for non-autoregressive models.These approaches therefore do not directly address autoregressive LLM inference.
- Token Dropping and KV Cache Compression: Many existing token-compression methods require additional retraining, limiting their suitability for autoregressive LLMs such as ChatGPT and Llama.The related-work discussion identifies this as a gap motivating KV-cache pruning.
- Underlying Structure of Attention: Research on self-attention has shown that heads can have interpretable roles, differ in performance impact, and vary in importance across tasks.These findings support analyzing attention heads rather than treating them uniformly.
3 ADAPTIVE KV CACHE COMPRESSION
FastGen profiles attention during prompt encoding, selects a compression policy per attention head, and manages the cache during token generation. Its policy set combines special-token, punctuation, frequent-token, local, hybrid, and full-cache strategies under an attention-recovery objective.
- Generative inference: Generative inference consists of prompt encoding followed by token generation, during which the KV cache grows as new token vectors are appended.The cache stores previously computed key/value vectors to avoid recomputation.
- FastGen pipeline: FastGen profiles attention heads during prompt encoding, then applies the selected compression strategy while managing cache updates during token generation.The two phases are profiling and adaptive cache management.
- Model profiling: For each compression policy, FastGen chooses the minimum-memory-cost policy that recovers a predefined fraction T of the attention map.The policy operates on compressed key and value caches derived from the original K and V vectors.
- Model profiling: FastGen assumes that each head’s attention-map structure remains stable throughout generation, allowing policy selection from the encoded prompt alone.The paper reports theoretical and empirical support for using encoded prompts to capture full-context attention structures.
- KV cache compression policies: The compression policies retain special tokens, punctuation, or frequent tokens, evict distant contexts, or keep the full cache.Hybrid policies combine compressed caches, while Cfull denotes uncompressed storage.
4 DIVERSITY AND STABILITY OF ATTENTION STRUCTURES
Attention heads exhibit distinct structures across layers, and these patterns remain relatively stable during decoding, supporting head-specific KV-cache policies.
- 4.1 HEAD DISTINCTIVE ATTENTION STRUCTURE: Attention heads in different layers have vastly different structures, with initial and final layers containing more full-KV-cache heads than middle layers.
- 4.1 HEAD DISTINCTIVE ATTENTION STRUCTURE: Middle-layer attention heads often focus on special tokens, whose accumulated attention score exceeds 0.95 for these heads.
- 4.1 HEAD DISTINCTIVE ATTENTION STRUCTURE: These layer- and head-level differences make applying one KV cache uniformly across layers suboptimal and motivate detecting each head’s structure.
- 4.2 PROFILE TENDS TO BE CONSISTENT IN ONE SEQUENCE: For a given prompt, each head’s attention structure remains relatively consistent through the decoding process despite fluctuations in accumulated attention scores.
- 4.2 PROFILE TENDS TO BE CONSISTENT IN ONE SEQUENCE: Layer 33 Head 0 and Layer 23 Head 2 almost exclusively attend to special tokens, while Layer 23 Head 0 emphasizes locality and punctuation.
- 4.2 PROFILE TENDS TO BE CONSISTENT IN ONE SEQUENCE: Layer 23 Head 3 assigns more than 10% of attention to other tokens, making it suitable for an uncompressed full KV cache.
- 4.2 PROFILE TENDS TO BE CONSISTENT IN ONE SEQUENCE: A large portion of attention scores falls on special tokens across cases, supporting the greedy construction of hybrid policies.
5 EXPERIMENT
FastGen was evaluated across Llama 1 models, generative tasks, memory usage, latency, and profiling overhead. It consistently reduced KV-cache memory and latency while preserving generation quality, with benefits increasing for larger models and longer generations.
- Experimental Setup: FastGen was evaluated on Llama 1 models from 7B to 65B across code, math, knowledge, and reading-comprehension tasks.The benchmarks were HumanEval, GSM8k, NQ, and TQA, using standard task metrics and GPT4 pairwise evaluation for instruction-tuned variants.
- Performance–Memory Trade-off: FastGen consistently improved over non-adaptive KV-cache compression methods across the evaluated settings.For 30B models, FastGen with 50% cache compressed surpassed non-adaptive methods with 15% cache compressed.
- Performance–Memory Trade-off: 44.9% pruned ratio was achieved on Llama 1-65B at a 45% win rate, compared with 16.9% on Llama 1-7B.The reported pruned ratio increased with model size while preserving the same model quality.
- Memory Footprint Reduction: ∼40% memory reduction was achieved for Llama 1-65B at a win rate over 45%, with ∼30% for 30B and ∼20% for 13B and 7B models.Memory was measured with batch size 16, sequence length 512, and fp16 model weights.
- End-to-End Latency Improvement: 55.0% latency reduction over HF was achieved at generation length 16k, compared with 16.04% at length 512 for batch size 1.Relative speedup increased with generation length, and FastGen also outperformed DeepSpeed across settings.
- Profiling Cost: Profiling accounted for at most 0.35% of total generation duration and fell to 0.07% at generation length 1024.The extra memory required by the Cfrequent strategy was 0.78% of KV-cache storage.
6 CONCLUSION
FastGen improves LLM inference efficiency through lightweight model profiling and adaptive key-value caching while maintaining no visible quality loss. Future work includes combining it with other compression methods and efficient attention architectures.
- Conclusion: FastGen significantly improves LLM inference efficiency with lightweight model profiling and adaptive key-value caching, without visible quality loss.The authors identify combinations with quantization, distillation, and grouped-query attention as future exploration areas.
A.1 ABLATIONS
The ablations show that adaptive policy combinations and their introduction order materially affect the trade-off between KV cache compression and generation quality.
- Policy effects: Removing Cfrequent and Cspecial causes win-rate drops of 3.67% and 2.11%, respectively, under the full policy set.Cfrequent and Clocal also reduce more KV cache than the other policies, but standalone non-adaptive deployment is suboptimal.
- Experimental setup: The complementary-policy study evaluates win rate against full-cache settings on fine-tuned Llama 1-65B with AlpacaEval.The policy-order ablation uses the same model and benchmark context.
- Policy order: The current policy-introduction order achieves the highest win-rates and pruned ratios among the examined orders.All examined orders add Cspecial first, while alternative orders change the compression–quality trade-off.
- Policy order: Using Cfrequent →Clocal →Cpunct. improves KV cache compression ratio at the cost of generation quality.The targeted recovery ratio is fixed at 0.98 for the policy-order study.
A.2 SENSITIVITY STUDY.
FastGen’s hyper-parameters have little visible effect on generation quality but substantially affect the compression ratio.
- Hyper-parameter sensitivity: Generation quality remains above 45% winrate across all tested hyper-parameter settings.Changing hyper-parameters produces a relatively large change in compression ratio; lowering the frequency-policy ratio from 0.3 to 0.1 leads to more KV cache.