Source-linked AI summary
SGD-KV: Summarization Guided KV Cache Compression
Zeyu Liu, Woomin Song, Xuandi Fu, Sai Muralidhar Jayanthi, Vivek Govindan, Aram Galstyan, Sravan Babu Bodapati, Srikanth Ronanki
TL;DR
Long-context LLM inference is constrained by KV-cache memory growth and compression methods that overlook attention-head specialization. SGD-KV identifies summarization heads with a chunk-summarization diagnostic and allocates cache budget accordingly, achieving state-of-the-art long-context performance with up to 75% lower KV-cache memory usage.
Problem
KV-cache memory grows linearly with context length, while existing compression heuristics often overlook attention heads specialized for hierarchical information aggregation.
Method
SGD-KV uses a chunk-summarization diagnostic to identify summarization heads and allocates the compressible KV-cache budget according to their normalized summarization scores.
Results
SGD-KV achieves state-of-the-art performance across long-context benchmarks and reduces KV-cache memory usage by up to 75%.
Takeaways & Limitations
Summarization-guided, fine-grained head-level allocation provides an effective efficiency-accuracy strategy for long-context inference.
Takeaways & Limitations
Without the final query, performance degrades for SGD-KV, although a generic summarization prompt substantially mitigates the loss.
Abstract
from arXiv · showhide
Large language models (LLMs) face severe memory bottlenecks in long-context inference due to the linearly growing size of key-value (KV) caches. Existing KV cache compression techniques typically rely on simple heuristics, overlooking the distinct functional roles of different attention heads. We present SGD-KV (Summarization-Guided KV Cache Compression), a head-aware framework that leverages a novel chunk-summarization diagnostic task to systematically identify and prioritize attention heads specialized in hierarchical information aggregation. Experiments on Qwen2.5-7B-1M and Qwen3-32B across diverse long-context benchmarks demonstrate that SGD-KV achieves state-of-the-art performance with contexts up to 1M tokens, while reducing KV cache memory usage by up to 75%. Our findings show that strategically allocating the KV cache budget based on the summarization score distribution of attention heads yields a superior efficiency-accuracy trade-off for long-context inference.
1 Introduction
Million-token context inference is constrained by KV cache memory that grows linearly with context length. SGD-KV addresses this by identifying summarization-specialized attention heads and allocating cache budget accordingly.
- KV cache memory grows linearly with context length, creating a severe bottleneck for million-token LLM inference.
- SGD-KV identifies attention heads specialized in hierarchical information aggregation through a chunk-summarization diagnostic task.
- The framework uses summarization scores and a water-filling-inspired algorithm to prioritize heads during KV cache budget allocation.
2 Related Work
KV cache compression research has progressed from token eviction to head- and layer-aware allocation, but prior strategies largely use retrieval or attention-pattern heuristics. SGD-KV introduces summarization-based head specialization to capture hierarchical information synthesis.
- KV Cache Compression: Token-level methods discard cache entries using recency or cumulative-attention heuristics, without modeling functional roles.
- KV Cache Compression: PyramidKV and AdaKV allocate cache dynamically across layers and heads, but remain guided by quantitative attention patterns rather than semantic head roles.
- Attention Head Specialization: Prior head-specialization diagnostics identified retrieval, induction, and retrieval-reasoning heads through retrieval-centric tasks.
- Attention Head Specialization: SGD-KV introduces a summarization-based task to identify heads performing hierarchical information synthesis and applies that insight to KV compression.
3 Method
SGD-KV measures each attention head’s ability to connect generated chunk keywords to their source text, then uses the resulting scores to allocate KV cache budget. Higher-scoring heads retain more historical context, with token selection performed within each allocation.
- Chunk-Summarization Task Design: The chunk-summarization task concatenates documents, identifies semantic chunks, and extracts keywords for each chunk.
- Summarization Score Calculation: Each head receives an importance score measuring attention from generated keywords to corresponding source text within the correct chunk.
- Summarization Score Calculation: The final summarization score averages each head’s importance score across valid samples, using the strongest keyword-to-source attention link.
- KV Cache Budget Allocation: SGD-KV preserves sink and recent-token budgets, then allocates the remaining middle-sequence budget across heads in proportion to normalized summarization scores.
- KV Cache Budget Allocation: Heads identified as crucial for summarization receive larger budgets, after which cumulative-attention token selection retains important KV entries.
4 Experiments
Experiments evaluate SGD-KV on Qwen2.5-7B-Instruct-1M and Qwen3-32B across long-context retrieval, reasoning, and information-coverage benchmarks. The method consistently performs strongly, including 28.38 Avg. on ETHIC with 25% KV cache versus 28.53 Avg. for Full KV.
- Experimental Setup: SGD-KV is evaluated on Qwen2.5-7B-Instruct-1M and Qwen3-32B, representing non-reasoning small and reasoning-capable large models.
- MRCR: On MRCR, SGD-KV outperforms other SOTA head-level allocation methods on fine-tuned Qwen2.5-7B-1M and surpasses DuoAttention beyond 128K tokens.
- ETHIC: 28.38 Avg. with 25% KV cache on Qwen3-32B compares with 28.53 Avg. for the Full KV baseline on ETHIC.
- ETHIC: SGD-KV achieves SOTA performance on ETHIC for both evaluated models.
5 Ablation Study: Impact of Query on Token Selection
The ablation evaluates query-aware, query-unaware, and proxy-query token selection on MRCR, showing that query removal degrades performance while a summarization prompt mitigates the loss.
- Query Conditions: Token selection uses query-aware, query-unaware, and proxy-query conditions to test whether the final question is needed during compression.The default condition uses the last 128 tokens including the real query; query-unaware excludes it, while proxy-query uses a fixed summarization prompt.
- Results: Removing the final question causes significant performance degradation for both SGD-KV and HeadKV relative to query-aware performance.
- Results: A summarization prompt as proxy query substantially mitigates the performance loss when the final question is unavailable.The result supports using semantically salient summarization cues for token selection without the specific user query.
- Results: SGD-KV consistently outperforms HeadKV across all three query conditions.
6 Conclusion
The paper introduces summarization heads and integrates them into SGD-KV to improve long-context KV cache management. It reports state-of-the-art results, up to 75% memory reduction, and effective proxy-query token selection when the final question is unavailable.
- Conclusion: SGD-KV introduces summarization heads, identified with a chunk-summarization diagnostic task, to support hierarchical information synthesis.
- Conclusion: 75% maximum KV cache memory reduction accompanies state-of-the-art performance on complex, long-context benchmarks.
- Conclusion: A generic summarization prompt effectively serves as a proxy query for token selection when the final question is unavailable.
A.1 KV Cache Budget Allocation & Redistribution
The redistribution procedure caps over-allocated heads and reallocates excess budget toward high-scoring summarization heads while preserving per-head budget constraints.
- Budget Redistribution: The method uses a water-filling-inspired redistribution algorithm when high-importance heads exceed the middle-sequence budget.This situation can occur under a relatively large KV cache budget.
- Algorithm: The algorithm initializes per-head budgets, computes excess above M, caps over-allocated heads at M, and redistributes the excess.
- Algorithm: Excess budget is distributed to the highest-scoring heads that still have available capacity.
- Design Goal: The procedure prioritizes important summarization heads while maintaining budget constraints for effective long-context summarization.
A.2 Fine-tuning details
The appendix describes the fine-tuning data, recipe, baseline setup, and Figure 2’s IoU comparisons for summarization-head identification.
- Fine-tuning Dataset: The fine-tuning mixture contains 10K MRCR, 20K OpenAI GraphWalks, 25K BABILong, 10K Gutenberg, and 10K Llama-Nemotron samples.Gutenberg and Llama-Nemotron data are included for regularization to prevent overfitting.
- Fine-tuning Recipe: Full-parameter supervised fine-tuning uses Qwen2.5-7B-Instruct-1M, a 1.0 × 10^-5 initial learning rate, batch size 128, and two epochs.The recipe also uses cosine scheduling, warmup, FlashAttention, DeepSpeed stage 0, and Liger kernels.
- Baselines: Baselines include DuoAttention, AdaKV, and HeadKV, representing retrieval-versus-streaming classification and head-level KV cache allocation strategies.
- Configuration: Baseline configurations preserve sink and recent-token windows, with 1024-token settings for Qwen2.5-7B-1M and 128-token settings for Qwen3-32B.Allocations occur before repeating KV values, so reported savings build upon grouped-query attention.
- Head Identification: Figure 2 compares IoU scores across datasets and across different attention-head types.
A.3 Quantitatively Analysis of the Stability and Generalization of the Summarization Heads
SGD-KV’s summarization-head rankings are stable within datasets and generalize across summarization tasks, while retaining patterns distinct from retrieval-oriented heads.
- Scoring: Summarization heads remain continuously scored rather than being assigned by hard classification thresholds.This follows HeadKV’s approach of preserving importance scores for all heads.
- Stability: IoU exceeds 0.9 between disjoint Databricks Dolly subsets, indicating robust summarization-head identification within the dataset.The analysis ranks heads by importance and compares their top-k sets using Intersection over Union.
- Generalization: Cross-dataset comparisons among Dolly, DialogSum, and SAMSum maintain high IoU scores, especially for top-ranked heads.The reported comparisons indicate that summarization-head rankings generalize across datasets.
- Configuration: The final summarization-head configuration averages scores across six datasets, including DialogSum, SAMSum, CNN/Dailymail, XSum, Dolly, and WikiLingua.
- Head specialization: The top 20% of summarization, retrieval, and R2 heads overlap substantially, but the 20-60% range shows distinct preferences for summarization heads.This pattern indicates that summarization heads capture attention behavior beyond simple retrieval mechanisms.
A.4 Visualization of Different Types of Heads
The visualizations and variant experiments show that summarization-guided allocation overlaps with retrieval-aware allocation but emphasizes different layer-head patterns and score-guided budget distribution. Across cache budgets and configurations, SGD-KV generally outperforms competing allocation strategies, while low-scoring heads remain relevant at ultra-long contexts.
- Visualization of Different Types of Heads: Figure 3 maps head scores across 4 heads and 28 layers for DuoAttention, HeadKV, and SGD-KV.The rows represent head indices, while the columns represent layer indices.
- Visualization of Different Types of Heads: R2 and summarization heads show similar KV cache distributions, but SGD-KV allocates relatively more budget to earlier layers.The paper associates earlier layers with complex reasoning and summarization tasks.
- Performance under Different KV Cache Budget: SGD-KV consistently surpasses AdaKV and HeadKV across almost all KV cache budgets, with narrower margins below 15% or above 50%.The comparison is reported for MRCR accuracy under different cache budgets.
- Performance under Different KV Cache Budget: Above 50% cache budgets, SGD-KV and HeadKV outperform MInference particularly beyond 256K context, supporting head-level eviction over token-level approaches under adequate cache.The paper notes that this aligns with BABILong results under identical budget constraints.
- Head Configurations Variants: Reversing SGD-KV’s allocation performs significantly worse than the original method even with 3× KV cache, showing the importance of score-guided distribution.Averaging R2 and summarization scores produces performance between HeadKV and SGD-KV.
- Head Configurations Variants: Suppressing heads below the mean improves accuracy through 64K but degrades it beyond 128K, indicating low-scoring heads remain important for ultra-long contexts.
- Head Configurations Variants: For GQA, summarization scores must be aggregated from individual attention heads into one score per KV head.The Qwen2.5-7B-Instruct-1M example has 28 attention heads, 4 KV heads, and group size 7.
- Head Configurations Variants: All three GQA aggregation methods improve performance, with Ipt., Max achieving the best overall accuracy from 8K to 512K contexts.This result supports incorporating summarization scores into attention-score aggregation for GQA models.