Search papers, labs, and topics across Lattice.
This paper introduces SGD-KV, a head-aware framework for compressing key-value (KV) caches in large language models (LLMs) by utilizing a novel chunk-summarization diagnostic task to identify and prioritize attention heads that excel in hierarchical information aggregation. The approach significantly reduces KV cache memory usage by up to 75% while maintaining state-of-the-art performance on long-context benchmarks with contexts extending to 1M tokens. The results highlight the importance of strategically allocating KV cache resources based on the summarization capabilities of different attention heads, leading to improved efficiency in long-context inference.
By reallocating KV cache resources based on attention head specialization, SGD-KV cuts memory usage by 75% while handling contexts up to 1M tokens with state-of-the-art accuracy.
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.
