KVSink: Understanding and Enhancing the Preservation of Attention Sinks in KV Cache Quantization for LLMs

TL;DR

This paper investigates the mechanisms of attention sinks in large language models during inference, analyzes their impact on KV cache quantization, and introduces KVSink, a method that improves preservation of attention sinks, enhancing model efficiency and accuracy.

Contribution

It provides a detailed understanding of attention sinks in LLMs and proposes KVSink, a novel, low-overhead method for better preservation of attention sinks during KV cache quantization.

Findings

KVSink outperforms the Preserve-First-N strategy in preserving attention sinks.

Applying KVSink to KVQuant improves perplexity and reduces outliers.

Understanding of attention sink dynamics enhances quantization strategies.

Abstract

Key-Value (KV) cache quantization has become a widely adopted optimization technique for efficient large language models (LLMs) inference by reducing KV cache memory usage and mitigating memory-bound constraints. Recent studies have emphasized the importance of preserving the original precision of KVs for the first few tokens to ensure the protection of attention sinks. While this approach has proven effective in mitigating performance degradation, its underlying principles remain insufficiently understood. Moreover, it fails to address the recent discovery that attention sinks can emerge beyond the initial token positions. In this work, we elucidate the underlying mechanisms of attention sinks during inference by examining their role in the cross-layer evolution of extreme activation outliers. Additionally, we provide a comprehensive analysis of the interplay between attention sinks and KV cache quantization. Based on our enhanced understanding, we introduce \textit{\textbf{KVSink}}, a plug-and-play method that effectively predicts sink tokens with negligible overhead, enabling more thorough preservation. Extensive experiments demonstrate that KVSink outperforms the existing Preserve-First-N (PFN) strategy, offering more effective preservation of attention sinks during KV cache quantization. Moreover, when applied to the well-established KVQuant method, KVSink further improves perplexity (PPL) and reduces reliance on 16-bit numerical outliers.