FastKV：トークン選択的伝播を用いた高速な長文脈処理のためのKVキャッシュ圧縮

FastKV: KV Cache Compression for Fast Long-Context Processing with Token-Selective Propagation

February 3, 2025

著者: Dongwon Jo, Jiwon Song, Yulhwa Kim, Jae-Joon Kim

cs.AI

要旨

大規模言語モデル（LLMs）は長いコンテキストシーケンスを処理するのに優れていますが、コンテキスト情報を格納するためにかなりのキー値（KV）キャッシュが必要であり、これは計算効率とメモリ使用量に大きな負担をかける可能性があります。これまでのKVキャッシュの圧縮に関する努力は、主にメモリ要件を削減することに焦点を当てていましたが、レイテンシを向上させることには限界がありました。この問題に対処するために、私たちはFastKVを導入します。これは、長いコンテキストシーケンスのレイテンシを向上させるために設計されたKVキャッシュ圧縮メソッドです。処理速度を向上させつつ精度を維持するために、FastKVは新しいトークン選択伝播（TSP）アプローチを採用しています。これにより、LLMsの初期層で完全なコンテキスト情報を保持し、より深い層ではこの情報の一部のみを選択的に伝播させ、プリフィル段階でも同様です。さらに、FastKVにはグループ化クエリアテンション（GQA）に注意したKVキャッシュ圧縮が組み込まれており、GQAのメモリ効率と計算効率の両方の利点を活用しています。実験結果によると、FastKVは、最先端のKVキャッシュ圧縮方法であるHeadKVと比較して、最初のトークン到達時間（TTFT）とスループットにそれぞれ2.00倍と1.40倍の改善を達成しています。さらに、FastKVは、ベンチマークでの精度をベースラインと同等のレベルで維持しています。コードはhttps://github.com/dongwonjo/FastKV で入手可能です。

English

While large language models (LLMs) excel at handling long-context sequences, they require substantial key-value (KV) caches to store contextual information, which can heavily burden computational efficiency and memory usage. Previous efforts to compress these KV caches primarily focused on reducing memory demands but were limited in enhancing latency. To address this issue, we introduce FastKV, a KV cache compression method designed to enhance latency for long-context sequences. To enhance processing speeds while maintaining accuracy, FastKV adopts a novel Token-Selective Propagation (TSP) approach that retains the full context information in the initial layers of LLMs and selectively propagates only a portion of this information in deeper layers even in the prefill stage. Additionally, FastKV incorporates grouped-query attention (GQA)-aware KV cache compression to exploit the advantages of GQA in both memory and computational efficiency. Our experimental results show that FastKV achieves 2.00times and 1.40times improvements in time-to-first-token (TTFT) and throughput, respectively, compared to HeadKV, the state-of-the-art KV cache compression method. Moreover, FastKV successfully maintains accuracy on long-context benchmarks at levels comparable to the baselines. Our code is available at https://github.com/dongwonjo/FastKV.