FlashAttention-3: 非同期処理と低精度化でさらに高速かつ高精度なAttention

 (together.ai)
1 ポイント 投稿者 GN⁺ 2024-07-12 | 1件のコメント | WhatsAppで共有
  • TransformerのAttentionボトルネックは長いコンテキストを扱うLLMの学習・推論速度を制限してきたが、FlashAttention-3はHopper GPUの機能を活用してこれをさらに大きく削減しようとする新バージョン
  • 中核となるのは、Tensor CoreとTMAの非同期性を利用して計算とデータ移動を重ね合わせ、matmulとsoftmaxを交互に配置してGPUのアイドル時間を減らす方式
  • H100でFlashAttention-2は理論最大FLOPSの35%活用にとどまっていたが、FlashAttention-3はFP16で最大740 TFLOPS、H100理論値の75%に到達
  • FP8低精度ではLLM activationの外れ値が誤差を増幅するため、Hadamard transformベースのincoherent processingにより基準FP8 attention比で誤差を2.6倍低減
  • FlashAttention-3はGitHubで公開されており、Attentionコストが下がるほど、より長いコンテキスト処理やLLMの学習・推論効率改善の余地が広がる

FlashAttention-3の目標と性能

  • AttentionはTransformerの中核レイヤーだが、大規模言語モデルや長いコンテキストのアプリケーションでは主要なボトルネックになっている
  • FlashAttentionとFlashAttention-2は、GPUメモリの読み書きを減らしてAttentionを高速化するアプローチを切り開き、現在ではほとんどのライブラリがTransformerの学習と推論の高速化に利用している
  • このアプローチは、LLMのコンテキスト長がここ2年でGPT-3・OPTの2〜4KからGPT-4の128K、Llama 3の1Mまで伸びることに貢献してきた
  • FlashAttention-2はH100 GPUで理論最大FLOPSの35%活用にとどまっていたが、FlashAttention-3はHopper GPUの新機能を活用してこれを引き上げる
  • 公開されたFlashAttention-3の性能は次の通り
    • FP16でFlashAttention-2より1.5〜2.0倍高速
    • FP16で最大740 TFLOPS
    • H100理論最大FLOPSの75%活用
    • FP8でほぼ1.2 PFLOPS
    • 基準FP8 attention比で2.6倍小さい誤差

FlashAttention方式のおさらい

  • FlashAttentionはAttention計算の順序を再配置し、tilingと再計算を用いて高速化しつつ、sequence lengthに対するメモリ使用量をquadraticからlinearへ削減する
  • 入力ブロックをHBMからSRAMへロードし、そのブロックに対してAttentionを実行した後、出力をHBMに更新する
  • 大きな中間Attention行列をHBMに書き込まないためメモリの読み書きが減り、実行時間ベースで2〜4倍の高速化が可能
  • tilingとsoftmax rescalingを併用することで、ブロック単位で処理しながらも近似なしに正しい出力を得られる

Hopper GPUの機能: WGMMA、TMA、FP8

  • FlashAttention-2はAmpere A100 GPUで理論最大FLOPSの最大70%まで到達できるが、Hopper GPUの新機能を十分には活用できていない
  • FlashAttention-3はHopperの3つの機能を使う
    • WGMMA: Hopperの新しいTensor Coreを活用するwarpgroup matrix multiply-accumulate機能で、Ampereのmma.syncよりスループットが高い
    • TMA: global memoryとshared memoryの間のデータ転送を高速化する専用ハードウェアユニットで、index計算やout-of-bound predicationを処理してregister使用量を減らす
    • FP8: FP16に比べてTensor Coreスループットを2倍にできるが、より少ないビットで浮動小数点値を表現するため精度とのtradeoffがある
  • FlashAttention-3はNVIDIA CUTLASSの抽象化を用いてHopper機能を活用する
  • 新機能を使うようFlashAttentionを書き直すだけでも、FP16 forward pass性能はFlashAttention-2の約350 TFLOPSから540〜570 TFLOPSへ向上する

非同期性でGEMMとsoftmaxを重ねる

  • Attentionの主要演算はQ-K、P-V間のGEMMとsoftmax
  • 最新アクセラレータではmatmul以外の演算がmatmulよりはるかに遅く、softmaxのexponentialのようなspecial functionはfloating point multiply-addやmatrix multiply-addとは別ユニットで処理される
  • H100 SXM5はFP16 matrix multiplyで989 TFLOPSを提供するが、special functionスループットは3.9 TFLOPSで256倍低い
  • head dimension 128では、matmul FLOPSがexponentialより512倍多くても、exponentialがmatmulに対して時間の50%を占めうる
  • FP8ではmatmul FLOPSが2倍速くなる一方でexponential速度はそのままなので、matmulとsoftmaxを並列実行する重要性がさらに増す

  • warpgroup間のpingpong scheduling

    • GPUのwarp schedulerは、一部warpがGEMM結果を待つ間に別のwarpを実行し、ある程度のoverlapを自動的に実現する
    • FlashAttention-3はsynchronization barrierを使って、2つのwarpgroupのGEMMとsoftmaxを手動でよりうまく重ねる
    • warpgroup 1があるiterationのGEMM1と次のiterationのGEMM0を先に実行する
    • その後、warpgroup 2がGEMMを実行している間にwarpgroup 1がsoftmaxを処理する
    • このpingpong scheduleは、別のwarpgroupのGEMM実行時間の裏にsoftmaxを隠す方式
    • 実際のスケジューリングは図式ほど完全にきれいではないが、FP16 attention forward passでhead dimension 128、sequence length 8Kの条件で約570 TFLOPSを620 TFLOPSまで引き上げる

  • warpgroup内部のoverlap

    • 1つのwarpgroup内部でも、そのwarpgroupのGEMM実行中にsoftmaxの一部を実行できる
    • このpipeliningにより、FP16 attention forwardスループットは約620 TFLOPSから640〜660 TFLOPSへ向上する
    • ただしGEMM accumulatorとsoftmaxの入出力を同時に保持する必要があるため、register pressureが増える
    • 全体としては、この手法は有利なtradeoffをもたらす

FP8低精度とincoherent processing

  • LLM activationには、ほかのfeatureよりmagnitudeがはるかに大きいoutlierが存在することがある
  • outlierは量子化を難しくし、quantization errorを大きく増やす
  • FlashAttention-3は、QuIPのような量子化文献で使われるincoherent processingを活用する
  • queryとkeyにrandom orthogonal matrixを掛けてoutlierを分散させ、量子化誤差を減らす
  • 実装ではrandom signを含むHadamard transformを用いる
    • head dimensionをdとすると、O(d²)ではなく**O(d log d)**時間でattention headごとに実行できる
    • Hadamard transformはmemory-bandwidth boundなので、rotary embeddingのように同じくmemory-bandwidth boundな前段演算とfuseすれば追加コストなしで処理できる
  • Q、K、Vを標準正規分布から生成し、entryの0.1%に大きなmagnitudeを与えてoutlierをシミュレーションした実験では、incoherent processingは量子化誤差を2.6倍低減した

ベンチマークと公開状況

  • FlashAttention-3はFlashAttention-2だけでなく、Hopper GPUの新しいハードウェア機能をすでに使っているTritonおよびcuDNN実装とも比較されている
  • FP16ではFlashAttention-2比で約1.6〜1.8倍の高速化を示す
  • FP8ではほぼ1.2 PFLOPSに到達する
  • FlashAttention-3 GitHubリポジトリが公開された
  • 論文も同じflash-attentionリポジトリで確認できる

残る最適化と今後の統合

  • 論文には、ブログで扱われた内容以外にもvariable length sequence、persistent kernel、FP8向けin-kernel transposeなどの最適化が含まれている
  • 実行ハードウェアに合わせてアルゴリズムを設計すれば、大きな効率向上や長いコンテキストのような新しいモデル機能を切り開ける
  • 今後の作業には、LLM inferenceの最適化や他のハードウェアアーキテクチャへの手法の一般化が含まれる
  • FlashAttention-3は今後のPyTorchリリースに統合されることが期待されている

関連記事

1件のコメント

 
GN⁺ 2024-07-12
Hacker News の意見

  • コードのコメントを見ると、Tri Dao は Hopper/H100 発表直後の 2022 年 4 月から FA3 に取り組んでいたようです
    コードが今日公開されるまで 2 年以上かかった点は少し興味深く、おそらくより良い解法が準備中だからかもしれません
    最近の Tri の論文履歴は、SSM や Mamba 系の構造に傾いています。FlashAttention はシーケンス長に対して 二次時間計算量を持ちますが、最新のアルゴリズムは準二次なので、同じ計算をより効率よく行うというレベルではなく、計算量そのものを大幅に減らします
    Dao と Gu は今年の長い論文で、Mamba/SSM も Transformer が恩恵を受けるのと同じハードウェアプリミティブで高速化しやすいように定式化できることを示しました

    • 強い指数時間仮説（SETH） が証明または反証されるまでは、二次コストが必要になるか、何かを諦める必要があります。結局のところ全探索のコストです
      SETH を証明または反証すれば P 対 NP 問題も解決されるため、近いうちに実現すると期待するのは難しいでしょう
      重要なのは、特定のユースケースがそのコストを負担できるかどうかです

  • FlashAttention アルゴリズムがハードウェアにどの程度結び付いているのか気になります
    たとえば今回の発表では H100 GPU の 非同期機能を活用するとされていますが、そうすると H シリーズではないカードではその速度向上を得られないという意味に見えます
    また実際の FlashAttention ライブラリには CUDA が必要ですが、アルゴリズム自体は Metal に移植されているようです[^0]。アルゴリズムが純粋関数に近いものなら、どの GPU/機械学習フレームワークでも実装できるべきではないかと思います
    [0]: https://github.com/philipturner/metal-flash-attention

    • 良い回答は多いですが、短く言えば「実務的にはかなり」ハードウェアに結び付いています。以下はよい例です

      https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
      Karpathy の nanoGPT は torch.nn.functional.scaled_dot_product_attention の有無を確認して FlashAttention を呼び出します
      https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
      ドキュメントを見ると、実際には多くの場合 FA2 を呼び出したくなるもので、FA2 はデバイスのカーネルを最適化して三角行列の Softmax 演算を分割し、不要な浮動小数点のバッチを GPU と CPU の間で往復させる処理を減らします
      https://arxiv.org/pdf/2307.08691
      FA2 論文は、ほぼ全面的にそれが実行されるハードウェアの観点から記述されています

    • FlashAttention のアルゴリズム上の改善は主に、アテンションの Softmax 部分を分割して結合することであり、それ自体がまったく新しいアイデアというわけではありません。圧倒的な貢献は、その方法と細部を Nvidia ハードウェア上で効率的に実装した点にあります
    • もともとの FlashAttention はハードウェア依存性がほとんどありませんでした
      最新版は抽象化レベルによって異なります。ThunderKittens[0] は記事で述べられているものと似た形で、FA2 比 1.3〜2 倍程度の速度向上を提供しつつ、GPU 全般に比較的広く適用できます
      新しいハードウェアごとに、追加性能を引き出せるハードウェア専用機能がある場合があります。通常、ベンダーは自社を先行させる機能を採用しますが、CUDA でもすでにそうであるように、API とライブラリは断片化します
      [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
    • 概念的には少し、実務上の実装という観点では大きく結び付いています。標準的な Python 実装も内部では特定のハードウェア向けに カーネルをコンパイルします
    • 実務面から付け加えると、AMD ハードウェアにはまだ flash-attention-2 のきちんとした実装が不足しています。ROCm は徐々に使えるようになってきていますが、CUDA と比較できる水準にはまだ達していません

  • コンパイラ方面の方々に聞きたい。コンパイラがFlashAttentionのような最適化を自力で見つけ出せる可能性はあるのだろうか？ TVMとtinygradはその方向に進んでいるように見えるが、可能だとは信じにくい

    • 理論的には可能。数学の代数的性質のおかげで大幅な並べ替えが可能で、そこに比較的定型的な多面体ループタイリングを加えればよい
      ただしコストが大きいので、その探索結果をキャッシュする必要がある
      e-graph最適化はこの領域によく合っていそう。ただ、最適化パスの処理方式に大きなパラダイム転換が必要で、一部のニッチなツール以外ではほとんど配布されていない。例えば従来のコールグラフとは相性が悪く、基本ブロックやforループの外側／間をまたいでe-graphを展開するには制御フローを大きく変える必要があり、breakやreturnもサポートされない
    • 極めて難しいが、不可能ではない問題に見える
      コンパイラ最適化の最新水準が、データ配置とプロセッサ使用率の最大化という面でどこまで来ているのかはよく分からない
      以前、最適化に関する動画で、小さな最適化は速度を上げはしたものの、その最適化や、さらにはランダムな変更が生んだメモリ配置の違いによる速度変動に比べると影響はごく小さい、という内容を見たことがある
      その発表はノイズの中からシグナルを見分けることにより焦点を当てていたが、そのノイズ自体が、コンパイラがここで述べた問題よりはるかに単純な形でさえうまく扱えていない痕跡だ
      キャッシュとアクセスパターンが速度に影響するCPU・メモリ構造だけでも複雑なのに、ここにGPU構造まで加えると、かなり未開拓の領域だと思う
      いつかは可能になるかもしれない。AI分野なので、十分に賢いAIがこれをできるのかという問いも出てくるが、「十分に」の基準次第だ
      AIモデルの非常に高水準なテストとして、microgradのようなものを与えて、同じインターフェイスを維持しながらtorchより速いものを作れと指示する、という形を思い浮かべられる。まだその近くにも到達していないが、可能になれば面白いだろう
    • そうではないと思う。別のアルゴリズムのように考えるべきだ。数学だけを考慮するのではなく、ハードウェアの形に合わせてアルゴリズムを設計するということだ
      TVMは納得できる。厳密には別のことをしているが、かなり近い領域だ
      ただ、tinygradについてはなぜそう感じたのか分からない
    • https://github.com/uwplse/tensat
    • Pythonのようなラッパー言語から高水準演算子を呼び出すのはかなり厄介だ

  • これをROCm / AMD MI300xに移植したい人がいれば、hello@hotaisle.xyz まで連絡してほしい。スパムは絶対に送らない
    この作業のための計算時間を提供できる

    • AMDアクセラレータのサーバー会社なんですね！すばらしい取り組みで、誰かが引き受けてくれるといいですね :)
    • 失礼なつもりはないが、この提案の意図が気になる。ハードウェアへのアクセス権だけを受け取って、誰がこの移植を無料でやってくれるのだろう？その人にはどんなメリットがあるのか？

  • FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
    FA3は3090や4090のようなコンシューマ向けGPUではどの程度の性能が出るのか？

  • TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
    私の理解では、TMAがレジスタを減らしてくれるのも確かだが、より重要なのはアドレス生成をハードウェアに処理させられる点だ。周辺の演算が速くなるほど、アドレス生成がボトルネックになり得る

  • これはAI全体で最も重要な改善の一つだ。同じハードウェアでより多くのことをより速く使えるようにし、ほとんどのAIユーザーにほぼトレードオフなしの利益をもたらす

    • H100を持っているユーザーにとってはそうだろう

  • FlashAttentionが可変マスキングを使うと、使わない場合よりなぜ5倍ほど遅いのか気になる。優れたマスキング対応が不足していると、最適化効果がほとんど消えてしまう

    • そのベンチマークはどこで見ていますか？

  • 専門家がいくつかの質問に答えてくれるとうれしいです :)
    FlashAttentionはLLMのアテンション演算をそのまま置き換えられるドロップインなのでしょうか？ 「アテンション」演算が使われる場所ならどこでも使えるのか、それともFAを使うようにLLMを別途学習する必要があるのでしょうか？
    FAはGQAやスライディングウィンドウ・アテンションのような戦略とどういう関係にありますか？ 互いに直交する概念なのでしょうか、それとも戦略ごとに別のFA実装が必要なのでしょうか？
    最近llama.cppがFlashAttention対応を追加しましたが、これはFlashAttentionが提供するCUDAカーネルのようなものを使い始めたという意味なのでしょうか？
    最後に、この記事ではFlashAttentionとTritonを比較しています。Tritonは抽象化レイヤーのようなものではないのですか？ FAをTritonで実装することはできないのでしょうか？ 「FlashAttention対Triton」という表現がよく理解できません。

      1. ほぼその通りです。数学的には等価です。ソフトウェア上の問題は依存関係のバージョン管理やメモリ内のデータ形式のようなものだけで、FlashAttention 2はすでにHuggingFaceや多くの人気ライブラリに入っています。FlashAttention 3も近いうちに入る可能性が高いですが、実行するにはH100 GPUが必要です。
      2. FlashAttention 2は以前のバージョンアップデートでGQA対応を追加しました:
        https://github.com/Dao-AILab/flash-attention
      3. ここでは、純粋なCUDA C++で書かれたこのFlashAttention実装と、Tritonで書かれた似たアルゴリズムのTriton実装を比較しているのです: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
    • FlashAttentionはLLMのアテンション演算をドロップインで置き換えられます。
      FlashAttentionはアテンションのSoftmax(QK^T)V部分を計算する方法で、GQAはQ、K、V行列を計算する方法です。スライディングウィンドウ・アテンションについてはあまり確信はありませんが、どのクエリがどのキーを見られるかを制御するアテンションマスクを変える方式です。
      llama.cppは使ったことがありませんが、CUDAカーネルを取り込んで使い始めたという説明でおおむね合っているように見えます。
      最後の質問は、以前にTritonで書かれたFlashAttention実装のことを指しています。

  • この記事でsigmoidのような演算が非常に遅いとあったので気になりました。
    最近のLLMはSiLU、Swish、SOLUのようにsigmoidやSoftmaxを含む活性化関数を多く使っています。
    ReLUは性能低下が少ないのでしょうか？ もしそうなら、昔ながらの良いReLUに戻るほうがよい可能性もあるのでしょうか？

    • ReLUは文字どおり、ある点で0に切られる線形関数なので、指数関数を含むものより計算量はずっと少ないです。ただ、そのように単純な活性化関数で競争力のある結果を得るのは難しいと思います。