Source-linked AI summary
SnapKV: LLM Knows What You are Looking for Before Generation
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, Deming Chen
TL;DR
Long prompts make KV caching costly in memory and decoding efficiency, especially because existing approaches often overlook prompt-cache compression. SnapKV uses stable per-head attention patterns from an observation window to select important prompt positions and compress the cache without fine-tuning. It reports comparable accuracy to full KV caching across diverse long-sequence evaluations while improving decoding speed and memory efficiency, with further testing reaching 380K-token contexts on one A100-80GB GPU.
Problem
Growing prompt lengths make KV caches increasingly costly in decoding time and memory, while prior methods often overlook compressing prompt KVs.
Method
SnapKV uses a fine-tuning-free observation-window pattern to identify critical per-head prompt attention features and compress the corresponding KV cache.
Results
SnapKV achieves comparable accuracy to full KV caching across diverse long-sequence datasets while improving decoding speed and memory efficiency.
Takeaways & Limitations
SnapKV maintains a constant prompt KV-cache size during generation, supporting more efficient long-context inference and processing up to 380K tokens on one A100-80GB GPU.
Takeaways & Limitations
SnapKV primarily targets generative KV caching and does not extend a model’s inherent long-context capability or cover prompt inference processing.
Abstract
from arXiv · showhide
Large Language Models (LLMs) have made remarkable progress in processing extensive contexts, with the Key-Value (KV) cache playing a vital role in enhancing their performance. However, the growth of the KV cache in response to increasing input length poses challenges to memory and time efficiency. To address this problem, this paper introduces SnapKV, an innovative and fine-tuning-free approach that efficiently minimizes KV cache size while still delivering comparable performance in real-world applications. We discover that each attention head in the model consistently focuses on specific prompt attention features during generation. Meanwhile, this robust pattern can be obtained from an 'observation' window located at the end of the prompts. Drawing on this insight, SnapKV automatically compresses KV caches by selecting clustered important KV positions for each attention head. Our approach significantly reduces the growing computational overhead and memory footprint when processing long input sequences. Specifically, SnapKV achieves a consistent decoding speed with a 3.6x increase in generation speed and an 8.2x enhancement in memory efficiency compared to the baseline when processing inputs of 16K tokens. At the same time, it maintains comparable performance to the baseline models across 16 long sequence datasets. Moreover, SnapKV can process up to 380K context tokens on a single A100-80GB GPU using HuggingFace implementation with minor changes, exhibiting only a negligible accuracy drop in the Needle-in-a-Haystack test. Further comprehensive studies suggest SnapKV's potential for practical applications.
1 Introduction
Long-context prompts make KV caching increasingly costly in decoding latency and memory, while existing approaches often overlook prompt-cache compression. SnapKV identifies stable prompt attention patterns before generation and compresses the cache without fine-tuning, preserving comparable accuracy while improving efficiency.
- Longer prompts increase per-step decoding latency linearly and require more memory because attention repeatedly accesses past KVs.
- SnapKV forms a new cache by combining selected per-head prefix features with the observation-window features for generation.
- Existing KV-cache methods mainly compress entries added during decoding, leaving large prompt caches as a practical memory and latency bottleneck.
- SnapKV evaluates attention-allocation questions experimentally, including whether input-token patterns are consistent and identifiable before generation.
- SnapKV identifies critical, stable attention features before generation and compresses the corresponding KV cache without fine-tuning or major model modification.
- Across diverse LLMs and long-sequence datasets, SnapKV reports comparable accuracy to full KV caching alongside improved decoding speed and memory efficiency.
2 Related Works
Prior KV-cache compression methods selectively retain or drop entries, but commonly focus on decoding-time additions rather than the large prompt cache. The related work therefore leaves long-prompt retrieval and efficiency challenges insufficiently addressed.
- StreamLLM retains recent tokens and attention sinks, which can discard information carried by middle tokens.
- H2O drops KVs during generation using cumulative-attention scores but does not compress prompt KVs.
- FastGen extends KV compression with a dual-phase algorithm that includes multiple KV-cache components.
- Existing methods may reduce decoding-stage cache size while leaving long-prompt understanding, memory overhead, and accurate information retrieval unresolved.
3 Observations
The observations show that prompt attention patterns can be identified before generation and remain stable as generation proceeds. In particular, the final prompt window closely matches the attention allocation used during subsequent generation.
- The study analyzes attention allocation in Query-Key matrices using filtered Ultrachat sequences with long prompts and responses.
- Attention features are evaluated across 128-token prompt windows by comparing high-average-attention positions with those used during generation.
- The last prompt window recognizes a highly similar attention-allocation pattern to the pattern used during actual generation.
- Features identified as crucial in the last prompt window maintain high overlap across four subsequent generation windows.
4 SnapKV
SnapKV compresses prompt KV caches by using an observation window to identify important prefix attention features, clustering selected positions, and retaining them with the full observation window. Its context-aware selection is evaluated across long-document QA settings, where instruction content changes selected features but instruction position still yields consistently high hit rates.
- SnapKV maintains a constant amount of prompt KVs during generation to reduce long-context serving time.
- The observation window is the prompt’s last segment, while the prefix is the preceding input used for attention analysis.
- Voting aggregates observation-window attention weights across heads to identify significant prefix positions, selecting top-k indices with k = ⌊p × Lprefix⌋.
- The hit rate measures selected important features against thresholded features during generation, with overlap quantifying alignment between voting and current attention.
- SnapKV clusters selected attention features, retains surrounding positions, and concatenates them with all observation-window features for the compressed cache.
- 4.2.1 Contextual Dependency of Patterns: Different instructions prioritize different prefix features, challenging static compression, whereas hit rates remain consistently high when questions appear before or after extensive context.
- Naively retaining only top-attention features can lose complementary details that induction heads use for copying, risking incomplete information.
5 Experiments
The experiments evaluate SnapKV on long-context retrieval, decoding speed, memory limits, and pooling across LWM-Text-Chat-1M and Mistral-7B-Instruct-v0.2. SnapKV preserves retrieval while reducing latency and extending usable input lengths.
- 5.1.1 Needle-in-a-Haystack: SnapKV retrieved the needle correctly before 140K tokens with only a small accuracy drop afterward.The original implementation encountered an out-of-memory error at 33K input tokens.
- 5.1.1 Needle-in-a-Haystack: 380K tokens was the maximum document length tested on a single A100-80GB GPU for Needle-in-a-Haystack retrieval.SnapKV used a 1024-token prompt KV cache, maximum pooling kernel size 5, and observation window size 16.
- 5.1.2 Decoding Speed and Memory Bound: 3.6x speedup was achieved at 16K input tokens and batch size 2, with SnapKV decoding below 40 ms compared with over 100 ms for baseline.Baseline latency increased linearly with input length, whereas SnapKV maintained a constant decoding speed.
- 5.2 Ablation Study of Effectiveness of Pooling: Pooling enabled correct LongEval-Lines retrieval before 16K tokens and significantly outperformed the no-pooling variant.The ablation used noisy key-value lines and evaluated ground-truth positions from 5K to 30K tokens.
- 5.1.2 Decoding Speed and Memory Bound: 8.2x longer input capacity was achieved at batch size 2, extending the limit from beyond 16K baseline tokens to 131K with SnapKV.The baseline encountered an out-of-memory error beyond 16K input tokens.
5.2 Ablation Study of Effectiveness of Pooling
The pooling ablation examines how clustering affects retrieval in noisy LongEval-Lines contexts. Pooling improves retrieval, while naive KV compression can produce partially correct results by disrupting surrounding-token context.
- 5.2 Ablation Study of Effectiveness of Pooling: Pooling significantly improves LongEval-Lines retrieval compared with omitting pooling, enabling correct values before 16K tokens.The benchmark uses randomly generated adjective-noun keys and five-digit values embedded in noisy lines.
- 5.2 Ablation Study of Effectiveness of Pooling: LongEval-Lines is more difficult than Needle-in-a-Haystack because relevant key-value information is embedded in noisy contexts with the same format.
- 5.2 Ablation Study of Effectiveness of Pooling: Naive KV-cache compression can yield partially correct results because models copy tokens surrounding initial portions to preserve contextual integrity.This motivates pooling as a way to consolidate information through clustering.
- 5.2 Ablation Study of Effectiveness of Pooling: Max pooling and average pooling produced no significant performance difference throughout the experiments.
5.4 Experiments on Command-R
Command-R experiments test SnapKV in retrieval, generation, and end-to-end RAG settings under long contexts. SnapKV retains strong RAG performance while compressing KV caches across sequence lengths up to approximately 40,000 tokens.
- 5.4 Experiments on Command-R: Command-R evaluations used a maximum 4096-token KV cache, with pooling kernel size 13 and observation window size 64.These settings produced compression ratios between 2x and 32x depending on sequence length.
- 5.4.2 Retrieval Augmented Generation (RAG): 98.8% of Command-R’s baseline performance was retained for ground-truth document retrieval with 5–10x KV-cache compression.The benchmark used 20,000–40,000-token contexts and a 4096-token KV cache.
- 5.4.2 Retrieval Augmented Generation (RAG): SnapKV remained robust in generation quality and avoided the lost-in-the-middle pathology across relevant-document positions.The evaluation inserted relevant documents at the beginning, middle, and end of contexts.
- 5.4.2 Retrieval Augmented Generation (RAG): SnapKV improved over baseline when contexts contained close to 200 documents.The authors suggest compression may reduce noise from negative documents and focus attention on relevant information.
5.5 Case Study: Compatibility with Parallel Decoding
SnapKV complements parallel decoding by keeping prompt KV-cache size constant during generation, mitigating long-sequence slowdowns. In a 10k-token sequence, it outperformed both Medusa and native decoding.
- Compatibility with Parallel Decoding: SnapKV maintains a constant prompt KV-cache size during generation, addressing long-sequence computational bottlenecks in parallel decoding.The paper identifies query-key matrix multiplication tiling as one bottleneck when speculative decoding generates multiple tokens per step.
- Compatibility with Parallel Decoding: 1.3x speedup over Medusa and 2.2x over native decoding were achieved for sequences with 10k length.The baseline for native decoding is the HuggingFace implementation of naive decoding.
6 Discussions
SnapKV reduces KV-cache burdens during generation but does not extend a model’s inherent long-context capability or compress prompt inference.
- Discussions: SnapKV targets generative KV-cache compression and reduces computational and memory burdens from extensive prompts.Its method retrieves prompt information receiving consistent attention from specific heads during generation.
- Discussions: SnapKV cannot extend long-context capability when the underlying model already struggles with long contexts or performs poorly.This limitation concerns the model’s inherent context-processing capability rather than generation-stage cache management.
- Discussions: SnapKV does not cover prompt inference, limiting effectiveness when systems cannot handle extensively long prompts.The stated scope boundary applies before the generation stage.
A Discussion of Generation Time Speedup
SnapKV reduces generation-time overhead for long inputs by keeping decoding speed consistent as input length grows. Generation time dominates overall processing, while SnapKV balances prompting and generation time below 100k input length.
- Generation Time Speedup: SnapKV maintains consistent decoding speed regardless of input length, unlike the original model whose generation time increases with input length.The comparison used Mistral-7B-Instruct-v0.2 generating 512 tokens.
- Generation Time Speedup: Generation time dominates total processing time for LLM inputs, creating substantial overhead as input sequences lengthen.The paper documents prompting and generation stages separately to assess SnapKV’s effectiveness.
- Generation Time Speedup: SnapKV achieves balanced prompting and generation time for input lengths smaller than 100k.This comparison is presented in the prompting-time and generation-time breakdown.
B Visulization of the Generated Context
Figure 11 visualizes generated examples from three datasets, comparing ground truth with SnapKV at three prompt-token budgets and the full-KV-cache baseline.
- Visulization of the Generated Context: Generated examples from Samsum, Qasper, and HotpotQA are compared across ground truth, three SnapKV prompt-token settings, and the full-KV-cache baseline.The SnapKV settings use 1024, 2048, and 4096 prompt tokens.