GPU の無駄遣いをやめよう: FlashAttention がアップグレードされ、長いテキストの推論速度が 8 倍に向上

最近、ChatGPT や Llama のような大規模言語モデル (LLM) がかつてない注目を集めています。しかし、その運用コストは依然として非常に高いままです。 1 つの応答を生成するのにかかるコストはわずか 0.01 ドル (AWS の 8xA100 インスタンスで数秒で実行) ですが、数十億人のユーザーの需要を満たすためにスケーリングすると、コストはすぐに膨らみます。さらに、これらのユーザーは 1 日に複数回 LLM と対話する可能性があります。コードの自動生成など、一部のユースケースでは、新しい文字が入力されるたびに実行されるため、コストが高くなります。 LLM の採用が拡大するにつれて、生成時間のわずかな効率向上でも大きな影響が出るでしょう。

LLM 推論 (または「デコード」) は反復的なプロセスであり、トークンが 1 つずつ生成されます。 N 個のトークンを含む完全な文を生成するには、モデルを N 回順方向に通過させる必要があります。幸いなことに、以前に計算されたトークンをキャッシュすることができます。つまり、単一の生成ステップは、1 つの操作 (注意) を除いて、コンテキストの長さに依存しません。この操作により、コンテキストの長さが適切にスケーリングされなくなります。

LLM の重要な新たなユースケースの中には、より長いコンテキストを活用する必要があるものもあります。より長いコンテキスト ウィンドウがある場合にのみ、LLM は、ドキュメントを要約したり、ドキュメント内の質問に答えたりして、より長いドキュメントについて推論できます。さらに、より長い会話履歴を維持し、コードを 1 行も記述する前にコードベース全体を処理することもできます。たとえば、2022 年現在、ほとんどの LLM のコンテキスト長は最大 2k (GPT-3 など) ですが、現在では、一部のオープンソース LLM は 32k (Llama-2-32k など) まで拡張でき、一部のモデルは 100k (CodeLlama など) に達しています。このような状況では、注意操作が推論プロセスのかなりの部分の時間を占めます。

バッチ サイズの次元を拡大すると、コンテキストが比較的短い場合でも注意がボトルネックになる可能性があります。これは、バッチ次元が増加すると、読み取る必要があるメモリの量も増加するのに対し、モデルの残りの部分では、メモリ要件はモデルのサイズのみに依存するためです。

上記の問題を解決するために、FlashAttentionの著者であるTri Dao氏らは、「Flash-Decoding」と呼ばれる技術を提案しました。これにより、推論プロセス中の注意計算が大幅に高速化され、長いシーケンスの処理および生成速度が元の速度の8倍に向上しました。主なアイデアは、キーと値を可能な限り高速に並列にロードし、結果を個別に再スケーリングしてマージし、正しいアテンション出力を維持することです。

デコード中のマルチヘッドアテンション

デコード中、生成された各新しいトークンは、以前のすべてのトークンを参照して計算する必要があります: softmax (queries @ keys.transpose) @ values

この操作は、FlashAttention (最近の v1 および v2 バージョンを含む) を使用したトレーニング中に最適化されており、ボトルネックとなるのは中間結果 (Q@K^T など) の読み取りと書き込みのメモリ帯域幅です。ただし、ボトルネックが異なるため、これらの最適化は推論ケースに直接適用できません。トレーニング中、FlashAttention はバッチ サイズとクエリ長のディメンションを並列に処理します。推論中、クエリの長​​さは通常 1 です。つまり、バッチ サイズが GPU 上のストリーミング マルチプロセッサ (SM) の数よりも小さい場合 (たとえば、A100 には 108 個あります)、操作では GPU のごく一部しか使用されません。これは特に、GPU メモリに収まるように小さいバッチ サイズを必要とする長いコンテキストを処理する場合に当てはまります。バッチ サイズが 1 の場合、FlashAttention は GPU の 1% 未満を使用します。

FlashAttentionはクエリブロックとバッチサイズの間でのみ並列化され、デコード中にGPU全体を占有することはありません。

行列乗算プリミティブを使用してアテンション計算を実行することも可能であり、その場合、FlashAttention を使用する必要がなくなります。この場合、操作は GPU 全体を占有しますが、中間結果の書き込みと読み取りを行うカーネルが多数起動されるため、最適ではありません。

より速い注意力のデコード: フラッシュデコード

新しいメソッド Flash-Decoding は FlashAttention に基づいており、キー値シーケンスの長さという新しい並列次元を導入します。上記の 2 つの方法の利点を組み合わせたものです。 FlashAttention と同様に、グローバル メモリに追加データをほとんど保存しません。ただし、コンテキストが十分に長ければ、バッチ サイズが小さくても GPU を最大限に活用できます。

フラッシュデコードも、最終的な削減ステップを少し犠牲にして、キーと値の間で並列化されます。

フラッシュデコードは主に次の 3 つの作業ステップで構成されます。

• まず、キー/値を小さなチャンクに分割します。
• FlashAttentionを使用してクエリと各チャンクを並列化し、各行とチャンクに追加のスカラー値を書き込みます。アテンション値のlog-sum-exp
• 最後に、log-sum-exp を使用して各チャンクの寄与を調整し、すべてのチャンクを削減することで実際の出力が計算されます。

これはすべて、Attention/Softmax を反復できるため可能です。 Flash-Decoding では、ブロック内 (FlashAttention と同様) と、最終的な削減計算のためのブロック間の 2 つのレベルで使用されます。

実際には、キー/値ブロックは完全なキー/値テンソルのビューであるため、ステップ（1）ではGPU操作は行われません。次に、ステップ（2）と（3）をそれぞれ実行するために使用される2つの独立したカーネル関数があります。

CodeLlama 34B でベンチマーク

新しいアプローチを検証するために、研究者らは CodeLLaMa-34b のデコードスループットをベンチマークしました。このモデルは Llama 2 と同じアーキテクチャを持ち、一般的にその結果は多くの大規模言語モデルに適用できるはずです。研究者たちは、異なるシーケンス長（512kから64k）でのデコード速度をtok/sで測定し、注目度を計算するさまざまな方法を比較しました。

• Pytorch: 純粋な PyTorch プリミティブを使用してアテンション計算を実行します (FlashAttention を使用せずに)。
• フラッシュアテンションv2;
• FasterTransformer: FasterTransformer のアテンション カーネルを使用します。
• フラッシュデコード;

メモリとKVキャッシュからモデル全体を読み取るのに必要な時間を考慮した上限

非常に大きなシーケンスの場合、フラッシュデコードによりデコード速度が最大 8 倍高速化され、他の方法よりもスケーラビリティが大幅に向上します。

プロンプトが小さい場合、すべてのメソッドは同様に実行されます。ただし、シーケンス長が 512 から 64k に増加すると、フラッシュ デコードを除き、他の方法のスケーラビリティは非常に低くなります。このフラッシュデコードモード（バッチサイズは 1）では、シーケンス長を拡張しても生成速度に与える影響は小さくなります。

コンポーネントレベルのマイクロベンチマーク

研究者らはまた、さまざまなシーケンスの長さとバッチサイズを考慮して、f16 入力を使用した A100 でのマルチヘッド アテンションのマイクロベンチマークを実行しました。バッチ サイズを 1 に設定し、16 個の 128 次元クエリ ヘッドと 2 つのキー/値ヘッダー (グループ化されたクエリ アテンション) を使用します。これは、4 つの GPU で実行される CodeLLaMa-34b で使用される次元と一致します。

上記のマイクロベンチマークは、Multi-Head Attention の実行時間をマイクロ秒単位で示しています。フラッシュデコードは、シーケンスの長さが最大 ​​64k まで拡張された場合、ほぼ一定の実行時間を実現します。

これまでに測定されたエンドツーエンドの最大 8 倍の高速化は、アテンション自体が FlashAttention よりも最大 50 倍高速であるため可能です。フラッシュデコードは GPU を最大限に活用できるため、シーケンスの長さが 32k に達するまで、注目時間はほぼ一定です。

フラッシュデコードの使用

フラッシュデコードは次のリンクにあります:

• FlashAttention パッケージ、v2.2 以降: https://github.com/Dao-AILab/flash-attention/tree/main
• xFormers パッケージ (xformers.ops.memory_efficient_attention を検索)、0.0.22 以降: スケジューラは、問題のサイズに応じて、Flash-Decoding または FlashAttention メソッドを自動的に使用します。これらのメソッドがサポートされていない場合は、フラッシュ デコード アルゴリズムを実装する効率的な Triton カーネルにディスパッチできます。

LLaMa v2 / CodeLLaMa を使用した完全なデコード例は、FlashAttention リポジトリと xFormers リポジトリにあります。さらに、著者らは、高速で読みやすく、教育的で、変更しやすいように設計された、シンプルな LLaMa v1/v2 モデルの効率的なデコード コード例を提供しています。

参考リンク: https://princeton-nlp.github.io/flash-decoding/