Hydragen: High-Throughput LLM Inference with Shared Prefixes

要約

Transformer ベースの大規模言語モデル (LLM) は現在、数億のユーザーに導入されています。
LLM 推論は通常、数ショットのサンプルやチャットボット システム プロンプトなど、プレフィックスを共有するシーケンスのバッチに対して実行されます。
この大規模なバッチ設定でのデコードは、メモリーから大きなキー値 (KV) キャッシュを読み取り、バッチ内のすべてのシーケンスに対して非効率な行列とベクトルの積を計算するアテンション操作によってボトルネックになる可能性があります。
この作業では、共有プレフィックスを使用したアテンションのハードウェア認識型の正確な実装である Hydragen を紹介します。
Hydragen は、共有プレフィックスと一意のサフィックスに対して個別にアテンションを計算します。
この分解により、シーケンス全体でクエリをバッチ処理することで効率的なプレフィックス アテンションが可能になり、冗長なメモリ読み取りが削減され、ハードウェアに優しい行列乗算の使用が可能になります。
私たちの方法では、エンドツーエンドの LLM スループットを競合ベースラインと比較して最大 32 倍向上させることができ、バッチ サイズと共有プレフィックスの長さに応じて速度も向上します。
Hydragen では、非常に長い共有コンテキストの使用も可能になります。バッチ サイズが大きい場合、プレフィックス長を 1K から 16K トークンに増やすと、Hydragen のスループットは 15% 未満減少しますが、ベースラインのスループットは 90% 以上低下します。
Hydragen は、単純な接頭辞と接尾辞の分解を超えて一般化され、ツリーベースのプロンプト共有パターンに適用できるため、競技プログラミングの問題の推論時間をさらに 55% 削減できます。

要約(オリジナル)

Transformer-based large language models (LLMs) are now deployed to hundreds of millions of users. LLM inference is commonly performed on batches of sequences that share a prefix, such as few-shot examples or a chatbot system prompt. Decoding in this large-batch setting can be bottlenecked by the attention operation, which reads large key-value (KV) caches from memory and computes inefficient matrix-vector products for every sequence in the batch. In this work, we introduce Hydragen, a hardware-aware exact implementation of attention with shared prefixes. Hydragen computes attention over the shared prefix and unique suffixes separately. This decomposition enables efficient prefix attention by batching queries together across sequences, reducing redundant memory reads and enabling the use of hardware-friendly matrix multiplications. Our method can improve end-to-end LLM throughput by up to 32x against competitive baselines, with speedup growing with the batch size and shared prefix length. Hydragen also enables the use of very long shared contexts: with a high batch size, increasing the prefix length from 1K to 16K tokens decreases Hydragen throughput by less than 15%, while the throughput of baselines drops by over 90%. Hydragen generalizes beyond simple prefix-suffix decomposition and can be applied to tree-based prompt sharing patterns, allowing us to further reduce inference time on competitive programming problems by 55%.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google