HashEvict: A Pre-Attention KV Cache Eviction Strategy using Locality-Sensitive Hashing

要約

Transformer ベースの大規模言語モデル (LLM) は、キー/値 (KV) キャッシュを使用して、過去のトークンのキーと値の埋め込みを保存することで推論を大幅に高速化します。
ただし、このキャッシュは大量の GPU メモリを消費します。
この作業では、局所性依存ハッシュ (LSH) を使用して KV キャッシュを圧縮するアルゴリズムである HashEvict を紹介します。
HashEvict は、現在のクエリ トークンとコサインが異なるトークンをキャッシュ内で迅速に見つけます。
これは、現在のトークン クエリの 2 値化ガウス投影とキャッシュされたトークン キーの間のハミング距離を、埋め込み次元よりもはるかに小さい投影長で計算することによって実現されます。
これらの計算を容易にするために、GPU メモリ内に軽量のバイナリ構造を維持します。
アテンションを計算してトークン保持を決定する既存の圧縮戦略とは異なり、HashEvict はこれらの決定を事前アテンションで行うため、計算コストが削減されます。
さらに、HashEvict は動的です。デコードの各ステップで、現在のトークンのキーと値が、最も低いアテンション スコアを生成すると予想されるトークンの埋め込みを置き換えます。
HashEvict が、推論、複数選択、長いコンテキストの取得、および要約タスク全体で高いパフォーマンスを維持しながら、KV キャッシュを 30% ～ 70% 圧縮できることを実証します。

要約(オリジナル)

Transformer-based large language models (LLMs) use the key-value (KV) cache to significantly accelerate inference by storing the key and value embeddings of past tokens. However, this cache consumes significant GPU memory. In this work, we introduce HashEvict, an algorithm that uses locality-sensitive hashing (LSH) to compress the KV cache. HashEvict quickly locates tokens in the cache that are cosine dissimilar to the current query token. This is achieved by computing the Hamming distance between binarized Gaussian projections of the current token query and cached token keys, with a projection length much smaller than the embedding dimension. We maintain a lightweight binary structure in GPU memory to facilitate these calculations. Unlike existing compression strategies that compute attention to determine token retention, HashEvict makes these decisions pre-attention, thereby reducing computational costs. Additionally, HashEvict is dynamic – at every decoding step, the key and value of the current token replace the embeddings of a token expected to produce the lowest attention score. We demonstrate that HashEvict can compress the KV cache by 30%-70% while maintaining high performance across reasoning, multiple-choice, long-context retrieval and summarization tasks.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google