FlashAttention

この技術を一言で説明すると

GPUのメモリ間でデータを動かす回数を減らし、Attentionを効率よく計算する方法です。

なぜこの技術が必要なのか

長い入力ではAttentionの計算と中間結果が大きくなります。演算回数を減らす工夫だけでは、データを運ぶ時間が律速になって実時間の短縮につながらない場合があります。

従来は何が問題だったのか

単純な実装では、長さNに対してN×NのAttention行列を大きなメモリに書き出し、後の演算で読み直します。この読み書きと中間結果の保存が負担になります。

技術の仕組み

原論文の中心は、Q・K・Vを小さなブロックに分け、GPU上の高速なメモリを使って計算するタイル化です。softmaxを段階的に扱い、巨大な中間行列全体をHBMへ書き出すことを避けます。

標準の密なFlashAttentionは、Attentionそのものを近似する方式ではありません。ただし浮動小数点の演算順序が変われば数値の微差は起こり得ます。密なAttentionの理論的な演算量が線形になるわけでもありません。

具体的な使用例

対応する学習・推論基盤で、長い系列のAttention計算に使います。速度はGPU、dtype、入力形状、マスクなどに依存するので、実際の負荷で測定します。

類似技術との違い

メリット・デメリット

原論文の実験では効率化を示していますが、特定の高速化倍率をすべての環境に当てはめることはできません。未対応の形状では別のカーネルへ切り替わることもあり、設定名だけでは実際に使われたか判断できません。

実装方法

使用するフレームワークの対応GPU・精度・Attention条件を確認します。ウォームアップ後の実時間とピークメモリを測り、同じ入力で出力差も調べます。独自カーネルを実装するより、まず基盤が提供する対応実装で評価します。

最新動向

確認時のvLLM公式READMEは、FlashAttentionを含む複数の最適化Attentionカーネルを列挙しています。どの版が選ばれるかは環境に依存するため、本記事では原理を扱い、特定版を「常に最新・最速」とは表示しません。