Flash Attentionはローカル推論中のメモリトラフィックをどのように削減するのか？

エヴァ・ウォン

IceWhale author

エヴァ・ウォン は テクニカルライター であり ZimaSpaceの常駐ティンカーでもあります。 生涯のオタクであり、 ホームラボとオープンソースソフトウェアに情熱を持っています。彼女は複雑な技術的概念をわかりやすく、 実践的なガイドに翻訳することを専門としています。エヴァはセルフホスティングは楽しくあるべきで、怖がるものではないと信じています。彼女のチュートリアルを通じて、コミュニティが ハードウェアのセットアップを解明する手助けをしています。初めてのNAS構築からDockerコンテナの習得まで。

How Does Flash Attention Reduce Memory Traffic During Local Inference? - Zima Store Online

Flash Attentionは、正確なアテンションをタイル分割することでメモリトラフィックを削減し、中間スコアをGPUのHBMへ繰り返し書き込んだり読み戻したりする代わりに、チップ上で処理できるようにします。

ローカルAIサーバーでは、プロンプトが長くなるほどこの効果が重要になります。標準的なアテンションでは、中間処理の規模がますます大きくなるためです。Flash Attentionはモデルの知識を変更したり、KVキャッシュを削除したりするものではありません。GPUメモリ階層の中で最も遅い階層を通過するデータ量を減らすよう、アテンション計算を再構成します。

標準アテンションは大規模な中間テンソルをHBM上で移動させる

アテンションはクエリ、キー、バリューのテンソルから始まり、スコアを計算して正規化し、バリューと結合する必要があります。単純な実装では、これらの処理の間に大規模な中間行列を高帯域幅メモリへ書き出します。

初代FlashAttention論文は、アテンションの性能が演算量だけでなく、HBMとチップ上のSRAM間で行われる読み書きによっても制限されると論じています。スコア行列はトークンの組み合わせごとに大きくなるため、長いシーケンスほどこのトラフィックが増加します。

そのため、特にコンテキスト長が増えると、家庭用GPUは有用な行列演算よりもデータ移動に多くの時間を費やすことがあります。

Flash Attentionは問題をチップ上のワーキングセットにタイル分割する

Flash Attentionは、高速なチップ上メモリに収まるサイズのクエリ、キー、バリューのブロックを処理します。各タイルでは、中間データがチップ外へ出る前に複数のアテンション処理を実行します。

FlashAttention-3は、このファミリーの中核的な手法を、中間データのグローバルメモリトラフィックを回避するタイル分割戦略と説明しています。このアルゴリズムは、対応する定式化に対して正確なアテンションを引き続き生成します。

タイルサイズは、ハードウェアリソースやカーネル実装によって異なります。重要なのはデータ局所性です。読み込んだ値を演算ユニットの近くにある間に再利用します。

オンラインSoftmaxは完全なスコア行列の保存を不要にする

通常のSoftmaxでは、正規化の前に行内のすべてのスコアが必要に見えるため、実装では大規模な行列を materialize する必要があります。一方、Flash Attentionはタイルをまたいで正規化の統計情報を逐次的に保持します。

Tritonの融合アテンションチュートリアルでは、ブロック単位のSoftmax状態を使った融合アテンションを実演しています。これにより、部分的なスコアブロックを処理して破棄でき、すべての中間値をHBMへ書き戻す必要がなくなります。

ここでメモリトラフィックの削減と一時メモリの削減が結び付きます。アルゴリズムが計算を組み替えることで、大規模な中間オブジェクトをグローバルメモリ上に一度も存在させずに済むのです。

Featured

ZimaCube 2 パーソナルクラウドホームNAS

ZimaCube2

カーネル融合により、アテンション処理のより多くの工程を単一の実行経路にまとめられる

スコア計算、スケーリング、マスキング、正規化、バリューの集約は、個別のカーネルを起動して完全な中間データをグローバルメモリ経由で受け渡すのではなく、連携した1つのカーネルに融合できます。

PyTorchのscaled dot-product attentionインターフェースは、デバイス、データ型、形状、要求された機能が対応している場合、融合アテンション実装を選択できます。

その結果、プロンプト処理が高速化し、ピーク時の一時メモリ使用量が減る可能性があります。ただし、実際の向上幅は、ランタイムがFlash対応カーネルを実際にディスパッチするかどうかによって異なります。

IOの削減はKVキャッシュの排除と同じではない

自己回帰推論では、モデルまたはランタイムが別の破棄、圧縮、ウィンドウ化、オフロード方式を使用しない限り、後続トークンで利用するために、生成済みのキーとバリューを保持しておく必要があります。

AMDのROCmガイダンスでは、Flash Attentionがチップ上のストレージとHBM間のメモリ局所性を改善すると説明しています。この最適化が対象とするのは、アクティブな状態をアテンションがどのように読み取り計算するかであり、過去のKV状態が存在するかどうかではありません。

この違いを理解することで、よくある計画上の誤りを防げます。Flash AttentionはワーキングメモリとIOを削減できますが、保持されるKVキャッシュが増え続けるため、長いコンテキストを扱うサーバーの容量が不足することは依然としてあります。

高速経路が使われるかどうかはハードウェアと形状の対応状況で決まる

Flashカーネルは特化されています。ヘッド次元、データ型、アテンションマスク、GPUアーキテクチャ、ランタイムのバージョン、その他の機能によっては、別の実装へフォールバックすることがあります。

ZimaSpaceの記事「モデルパラメータを上回るアテンションメモリ」では、容量面を扱っています。Flash Attentionは、実行経路の最適化として別途測定する必要があります。

短いプロンプトと長いプロンプトでベンチマークを行い、選択されたバックエンドを確認したうえで、プリフィルのレイテンシ、割り当てメモリのピーク、メモリ帯域幅、出力トークンレートを記録してください。ZimaSpaceのローカルAIハードウェアガイドは、カーネル最適化では代替できない、より広範な容量面の基準を提供します。

NVIDIAのTransformer Engineは、より広範なカーネル選択スタックの一部として、最適化されたアテンションバックエンドとハードウェア依存の実行について説明しています。