RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

要約

トランスフォーマーベースの大規模言語モデル (LLM) はますます重要になっています。
ただし、アテンション計算の二次時間の複雑さのため、LLM をより長いコンテキストにスケーリングすると、推論レイテンシーが非常に遅くなり、キー値 (KV) ベクトルをキャッシュするための GPU メモリ消費量が多くなります。
この論文では、アテンションの計算を高速化し、GPU メモリ消費量を削減するためのトレーニング不要のアプローチである RetrievalAttend を提案します。
RetrievalAttend は、動的スパース性アテンション メカニズムを活用することにより、CPU メモリ内の KV ベクトルに近似最近傍検索 (ANNS) インデックスを使用し、生成中にベクトル検索で最も関連性の高いインデックスを取得することを提案します。
残念ながら、既製の ANNS インデックスは、アテンション メカニズムにおけるクエリ ベクトルとキー ベクトル間の分布外 (OOD) により、このような検索タスクには効果的でないことが多いことがわかります。
RetrievalAttend は、クエリ ベクトルの分布に適応できるアテンションを意識したベクトル検索アルゴリズムを設計することで、OOD の課題に対処します。
私たちの評価によると、RetrievalAttend は高いモデル精度を維持しながら、データの 1 ～ 3% にアクセスするだけで済みます。
これにより、GPU メモリのフットプリントが大幅に削減され、ロングコンテキスト LLM の推論コストが大幅に削減されます。
特に、RetrievalAttend では、8B パラメータを持つ LLM で 128K トークンを処理するために 1 台の NVIDIA RTX4090 (24GB) のみが必要で、0.188 秒で 1 つのトークンを生成できます。

要約(オリジナル)

Transformer-based Large Language Models (LLMs) have become increasingly important. However, due to the quadratic time complexity of attention computation, scaling LLMs to longer contexts incurs extremely slow inference latency and high GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to both accelerate attention computation and reduce GPU memory consumption. By leveraging the dynamic sparsity of attention mechanism, RetrievalAttention proposes to use approximate nearest neighbor search (ANNS) indexes for KV vectors in CPU memory and retrieves the most relevant ones with vector search during generation. Unfortunately, we observe that the off-the-shelf ANNS indexes are often ineffective for such retrieval tasks due to the out-of-distribution (OOD) between query vectors and key vectors in attention mechanism. RetrievalAttention addresses the OOD challenge by designing an attention-aware vector search algorithm that can adapt to the distribution of query vectors. Our evaluation shows that RetrievalAttention only needs to access 1–3% of data while maintaining high model accuracy. This leads to significant reduction in the inference cost of long-context LLMs with much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) for serving 128K tokens in LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google