A$^2$ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization

要約

長いコンテキストの大規模言語モデル（LLM）は、KVキャッシュの大きなメモリフットプリントと高いアクセスオーバーヘッドのために、効率的なサービングに大きな課題をもたらす。検索に基づくKVキャッシュ削減手法は、これらの課題を軽減することができる。典型的には、完全なKVキャッシュをCPUにオフロードし、推論中に必要なトークンをオンデマンドで検索する。しかし、これらの手法は依然として満足のいく精度劣化と余分な検索オーバーヘッドに悩まされている。これらの限界に対処するため、本論文では、検索に基づく新しいKVキャッシュ削減手法A$^2$ATSを提案する。A$^2$ATSは、ベクトル量子化技術をキー状態に適用することで、アテンションスコアの正確な近似を得ることを目的とし、それによりトップKトークンの効率的で正確な検索を可能にする。まず、Windowed Rotary Position Embeddingを提案し、位置埋め込み後にクエリとキー状態の位置依存性を切り離す。次に、注目スコア近似の目的を直接最適化する、クエリを考慮したベクトル量子化を提案する。最後に、KVキャッシュオフロードのための異種推論アーキテクチャを設計し、より大きなバッチサイズで長いコンテキストサービングを可能にする。実験の結果、A$^2$ATSは既存の手法と比較して、同等かそれ以下のオーバーヘッドで低い性能劣化を達成することができ、それによってロングコンテキストサービングのスループットが最大2.7倍$向上することが実証された。

要約(オリジナル)

Long context large language models (LLMs) pose significant challenges for efficient serving due to the large memory footprint and high access overhead of KV cache. Retrieval-based KV cache reduction methods can mitigate these challenges, typically by offloading the complete KV cache to CPU and retrieving necessary tokens on demand during inference. However, these methods still suffer from unsatisfactory accuracy degradation and extra retrieval overhead. To address these limitations, this paper proposes A$^2$ATS, a novel retrieval-based KV cache reduction method. A$^2$ATS aims to obtain an accurate approximation of attention scores by applying the vector quantization technique to key states, thereby enabling efficient and precise retrieval of the top-K tokens. First, we propose Windowed Rotary Position Embedding, which decouples the positional dependency from query and key states after position embedding. Then, we propose query-aware vector quantization that optimizes the objective of attention score approximation directly. Finally, we design the heterogeneous inference architecture for KV cache offloading, enabling long context serving with larger batch sizes. Experimental results demonstrate that A$^2$ATS can achieve a lower performance degradation with similar or lower overhead compared to existing methods, thereby increasing long context serving throughput by up to $2.7 \times$.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL