PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

要約

大規模言語モデル (LLM) は優れた理解能力を示していますが、推論中の GPU メモリ使用量の課題に直面しており、チャットボットなどのリアルタイム アプリケーションのスケーラビリティを妨げています。
推論を高速化するために、計算されたキーと値 (KV キャッシュ) を GPU メモリに保存します。
既存の方法では、KV キャッシュ圧縮を研究し、事前に計算された KV キャッシュをプルーニングすることでメモリを削減します。
ただし、層間の依存関係や事前計算における膨大なメモリ消費は無視されます。
これらの欠陥を調査するために、将来の世代に影響を与える重要なキーと値の数が層ごとに減少し、注意の重みの一貫性によってそれらを抽出できることがわかりました。
この発見に基づいて、重要なコンテキストを層ごとに保持することで KV キャッシュを圧縮する方法である PyramidInfer を提案します。
PyramidInfer は、パフォーマンスを犠牲にすることなく、より少ないキーと値を計算することでメモリを大幅に節約します。
実験結果では、PyramidInfer は、KV キャッシュの GPU メモリを 54% 以上削減し、Accelerate と比較してスループットが 2.2 倍向上することが示されています。

要約(オリジナル)

Large Language Models (LLMs) have shown remarkable comprehension abilities but face challenges in GPU memory usage during inference, hindering their scalability for real-time applications like chatbots. To accelerate inference, we store computed keys and values (KV cache) in the GPU memory. Existing methods study the KV cache compression to reduce memory by pruning the pre-computed KV cache. However, they neglect the inter-layer dependency between layers and huge memory consumption in pre-computation. To explore these deficiencies, we find that the number of crucial keys and values that influence future generations decreases layer by layer and we can extract them by the consistency in attention weights. Based on the findings, we propose PyramidInfer, a method that compresses the KV cache by layer-wise retaining crucial context. PyramidInfer saves significant memory by computing fewer keys and values without sacrificing performance. Experimental results show PyramidInfer improves 2.2x throughput compared to Accelerate with over 54% GPU memory reduction in KV cache.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google