RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

要約

トランスフォーマーベースの大規模言語モデル (LLM) は、さまざまなドメインでますます重要になっています。
ただし、アテンション操作の 2 次時間計算量は、推論レイテンシが非常に高く、キー値 (KV) ベクトルをキャッシュするための GPU メモリ消費量が大きいため、より長いコンテキストにスケーリングする際に重大な課題を引き起こします。
この論文では、アテンションの計算を高速化するためのトレーニング不要のアプローチである RetrievalAttend を提案します。
アテンションの動的スパース特性を活用するために、RetrievalAttend は CPU メモリ内の KV ベクトルに近似最近傍検索 (ANNS) インデックスを構築し、生成中にベクトル検索を通じて最も関連性の高いインデックスを取得します。
クエリ ベクトルとキー ベクトル間の分布外 (OOD) のため、既製の ANNS インデックスでは正確な取得のために O(N) (通常はすべてのキーの 30%) データをスキャンする必要があり、これを活用できません。
高いスパース性。
RetrievalAttention は、まず ANNS ベースのアテンションの OOD 課題を特定し、クエリに適応してデータの 1 ～ 3% のみにアクセスできるアテンション認識ベクトル検索アルゴリズムを介してそれに対処し、従ってサブリニアな時間計算量を実現します。
RetrievalAttend は、モデルの精度を維持しながら、GPU メモリ要件を大幅に削減して、ロングコンテキスト LLM の推論コストを大幅に削減します。
特に、RetrievalAttend は、8B パラメータを持つ LLM で 128K トークンを処理するために 16GB GPU メモリのみを必要とし、単一の NVIDIA RTX4090 (24GB) で 0.188 秒で 1 つのトークンを生成できます。

要約(オリジナル)

Transformer-based large Language Models (LLMs) become increasingly important in various domains. However, the quadratic time complexity of attention operation poses a significant challenge for scaling to longer contexts due to the extremely high inference latency and GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to accelerate attention computation. To leverage the dynamic sparse property of attention, RetrievalAttention builds approximate nearest neighbor search (ANNS) indexes upon KV vectors in CPU memory and retrieves the most relevant ones via vector search during generation. Due to the out-of-distribution (OOD) between query vectors and key vectors, off-the-shelf ANNS indexes still need to scan O(N) (usually 30% of all keys) data for accurate retrieval, which fails to exploit the high sparsity. RetrievalAttention first identifies the OOD challenge of ANNS-based attention, and addresses it via an attention-aware vector search algorithm that can adapt to queries and only access 1–3% of data, thus achieving a sub-linear time complexity. RetrievalAttention greatly reduces the inference cost of long-context LLM with much lower GPU memory requirements while maintaining the model accuracy. Especially, RetrievalAttention only needs 16GB GPU memory for serving 128K tokens in LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds on a single NVIDIA RTX4090 (24GB).

arxiv情報

提供元, 利用サービス

arxiv.jp, Google