今回の論文
今回取り上げるのは、Tri Dao による論文「FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning」です。2023年7月17日に arXiv で公開され、ICLR 2024 に採択された論文です。公開元は arXiv、URL は https://arxiv.org/abs/2307.08691、DOI は https://doi.org/10.48550/arXiv.2307.08691 です。研究分野は、Transformer、Attention高速化、GPUカーネル最適化、長文コンテキスト処理です。
この論文を選んだ理由は、FlashAttentionの次の改善点が「計算量を減らす」だけではなく、「同じAttention計算をGPUにどう割り当てるか」にあることを明確に示しているからです。LLMを長文対応にしたり、RAGやエージェントで大量のコンテキストを扱ったりする場合、Attentionの実装効率はプロダクトのコストと体感速度に直結します。FlashAttention-2は、モデル設計だけでなく、実行基盤の設計からAIアプリを考えるうえで参考になる論文です。
どんな技術か
FlashAttention-2は、TransformerのAttention計算をGPU上でさらに高速に実行するためのアルゴリズムと実装技術です。
もともとのFlashAttentionは、Attention行列を丸ごとGPUメモリに保存せず、ブロック単位で計算することで、メモリ読み書きを大きく減らしました。通常のAttentionは系列長に対してメモリ使用量が二次的に増えますが、FlashAttentionはタイル化と再計算を使い、正確なAttention結果を保ったままメモリ効率を改善します。
FlashAttention-2の主眼は、その先のボトルネックです。FlashAttentionはすでに高速でしたが、論文によるとA100 GPU上では理論ピークの25〜40%程度にとどまり、行列積に比べるとまだ余地がありました。FlashAttention-2は、GPUが得意な行列積の割合を増やし、1つのAttention head内でも並列化し、ワープ間の通信を減らすことで、AttentionカーネルをよりGEMMに近い効率へ近づけます。
言い換えると、FlashAttention-2は「Attentionを近似して軽くする技術」ではありません。計算結果は正確なAttentionのまま、GPUにとって実行しやすい形へ計算の順序と分担を作り直す技術です。
課題
Transformerを長いコンテキストに拡張すると、Attentionが大きなボトルネックになります。自己Attentionでは、各トークンが他のトークンと相互作用するため、系列長が長くなるほど計算量とメモリ使用量が急増します。
既存のFlashAttentionは、この問題に対して非常に有効でした。Attention行列を高帯域メモリに書き出さず、GPU内のより近いメモリ階層を活用することで、メモリ使用量を抑えながら高速化します。しかし、論文が指摘する限界は、FlashAttentionがまだGPUの計算資源を十分に使い切れていないことです。
難しいのは、GPUでは「FLOPsの数」だけで速度が決まらない点です。行列積はTensor Coreなどの専用ユニットで非常に速く実行できますが、softmaxのスケーリング、境界チェック、マスク処理、ワープ間の同期、共有メモリへの読み書きは相対的に高くつきます。A100ではFP16/BF16の行列積の理論性能に対して、非行列積のFP32演算性能はかなり低く、同じ1 FLOPでも実行コストが違います。
実際のAIシステムでは、この差は長文RAG、チャット履歴の長いアシスタント、コードベース理解、音声・動画の長系列処理、高解像度画像生成などで問題になります。モデルのコンテキスト長を伸ばしても、Attention実装が遅ければ、学習費用も推論レイテンシも増えてしまいます。FlashAttention-2は、長い系列を扱うための基盤技術として、この実行効率の課題を解こうとしています。
用語解説
- Attention
- Transformerで各トークンが他のトークンをどれだけ参照するかを計算する仕組みです。FlashAttention-2はAttentionの結果を近似せず、同じ計算をGPU上でより効率よく実行するため、この基本構造が前提になります。
- HBMとSRAM
- HBMはGPU上の大容量メモリ、SRAMはGPUチップ内の高速な小容量メモリです。FlashAttention系の重要な考え方は、巨大なAttention行列をHBMに何度も読み書きせず、SRAM上のブロック計算で済ませる点にあります。
- GEMM
- General Matrix Multiplication、つまり一般的な行列積です。現代GPUはGEMMを非常に高速に実行できるため、FlashAttention-2ではAttention計算のうちGEMMとして処理できる部分を増やし、非行列積演算を減らすことが重要になります。
- ワープ
- GPUで同時に命令を実行するスレッドのまとまりです。NVIDIA GPUでは一般に32スレッドが1ワープとして動きます。FlashAttention-2は、ワープ間でどのテンソルを分担するかを変えることで、共有メモリ経由の通信を減らします。
- Model FLOPs Utilization
- モデル学習で理論上必要なFLOPsに対して、GPUがどれだけ有効に使われているかを見る指標です。FlashAttention-2は、単体カーネルだけでなくGPT形式モデルの学習全体でも高い利用率を示しており、実運用の効果を判断するうえで重要です。
技術の仕組み
FlashAttention-2の仕組みは、大きく見ると「FlashAttentionのIO効率を保ったまま、GPUの計算資源をより埋める」ための再設計です。論文では、主に3つの工夫が示されています。
non-matmul FLOPsを減らす
1つ目の工夫は、行列積ではない演算を減らすことです。
Attention計算では、QK^T の行列積、softmax、softmax(QK^T)V の行列積が中心になります。行列積部分はGPUが得意ですが、softmaxのための最大値計算、指数関数、正規化、スケーリング、マスク処理などはTensor Coreに乗りにくく、相対的に高コストです。
FlashAttentionでは、ブロック単位でsoftmaxを正確に計算するためにオンラインsoftmaxのテクニックを使います。FlashAttention-2はこの計算を見直し、再スケーリング回数、境界チェック、causal maskに関わる処理を減らします。重要なのは、ここでAttentionを近似していない点です。出力を変えるのではなく、同じ出力を得るための途中計算をGPUに向いた形へ整理しています。
この工夫は、AIアプリ開発でも考え方として応用できます。高速化というとモデルを小さくする方向に目が行きがちですが、実際には「高コストな演算種類を減らし、ハードウェアが得意な演算へ寄せる」だけでも大きな差が出ます。
1つのAttention head内でも並列化する
2つ目の工夫は、並列化の粒度を細かくすることです。
初代FlashAttentionでは、主にバッチサイズとhead数の方向に並列化していました。つまり、1つのAttention headを1つのスレッドブロックが担当するような設計です。この方式は、バッチサイズやhead数が十分に大きい場合はうまく働きます。
しかし、長い系列を扱う場面では、バッチサイズを小さくせざるを得ないことがあります。GPUメモリに収めるためにバッチを小さくすると、起動できるスレッドブロック数が減り、Streaming Multiprocessorを十分に埋められません。A100のように多数のSMを持つGPUでは、これは低occupancy、つまり計算資源が空いている状態につながります。
FlashAttention-2は、1つのheadの中でも系列長方向に計算を分割し、複数のスレッドブロックで処理できるようにします。これにより、長文コンテキストでバッチが小さい場合でもGPU上の並列性を確保しやすくなります。
sliced-Kからsliced-Qへワーク分割を変える
3つ目の工夫は、スレッドブロック内のワープ間分担です。
初代FlashAttentionでは、各ワープがKとVを分割して持ち、Qを共有するsliced-K方式を使っていました。この方式では、各ワープが計算した中間結果を共有メモリに書き出し、同期し、足し合わせる必要があります。共有メモリはHBMより速いとはいえ、ワープ間通信や同期が増えるとカーネル全体の速度を押し下げます。
FlashAttention-2では、逆にQをワープ間で分割し、KとVを共有するsliced-Q方式を採用します。各ワープは自分が担当するQのスライスについて QK^T を計算し、そのまま同じVに掛けて出力の対応部分を作れます。出力スライスが分かれているため、ワープ間で中間結果を足し合わせる必要が少なくなります。
この変更は、一見すると低レベルな実装詳細に見えます。しかし、長文LLMのような大規模システムでは、このようなメモリアクセスと同期の削減が、実際のスループット差になります。モデルアーキテクチャの論文だけを読んでいると見落としがちですが、AIシステムでは「テンソルをどう分けるか」も重要な設計対象です。
対応範囲の拡大
FlashAttention-2は、head dimension 256までの対応、Multi-Query Attention、Grouped-Query Attentionにも対応します。これは、GPT-J、CodeGen系、Stable Diffusion 1.xのようなモデルや、KVキャッシュ削減を狙うLLM推論構成にとって実用上重要です。
特にGQAやMQAは、推論時のKVキャッシュを減らすためによく使われます。FlashAttention-2がこれらに対応することで、長文推論、バッチ推論、チャットサービスのスループット改善と組み合わせやすくなります。
実験と結果
論文では、Attentionカーネル単体のベンチマークと、GPT形式モデルの学習全体での速度評価が行われています。
Attentionカーネル単体の性能
A100 80GB SXM4 GPU上で、causal maskの有無、head dimension 64または128などの条件を変えて、Attentionのforwardとbackwardの速度を比較しています。
結果として、FlashAttention-2は初代FlashAttentionに対しておおむね約2倍高速になり、PyTorchの標準的なAttention実装と比べると最大で約9倍高速になったと報告されています。また、A100上で理論ピークの50〜73%に到達し、初代FlashAttentionの25〜40%から大きく改善しています。
ここで重要なのは、FlashAttention-2が近似Attentionではないことです。疎Attentionや局所Attentionのように参照範囲を削るのではなく、正確なAttentionをより効率よく実装しています。そのため、モデル品質を変えずに速度とメモリ効率を改善しやすい点が強みです。
GPT形式モデルの学習速度
論文では、GPT-3形式の1.3Bおよび2.7Bモデルを、2kまたは8kコンテキストで学習する設定も評価しています。
結果は、長いコンテキストほどFlashAttention-2の効果が大きく見えます。たとえばGPT-3 1.3Bの8kコンテキストでは、Megatron-LMベースラインが72 TFLOPs/s、FlashAttentionが170 TFLOPs/s、FlashAttention-2が220 TFLOPs/sです。GPT-3 2.7Bの8kコンテキストでは、ベースライン80 TFLOPs/s、FlashAttention175 TFLOPs/s、FlashAttention-2 225 TFLOPs/sと報告されています。
また、GPT形式モデルのend-to-end学習では、A100 1枚あたり最大225 TFLOPs/s、72%のModel FLOPs Utilizationに達しています。これは、カーネル単体だけで速いのではなく、モデル学習全体でも効果が出ることを示しています。
H100での結果
論文関連の技術解説では、H100 GPUでも同じ実装を走らせ、H100固有の新機能を特別に使わない状態で最大335 TFLOPs/sに到達したと紹介されています。これは、FlashAttention-2の設計が特定の1条件だけで効く小手先の最適化ではなく、GPU世代が進んでも有効な実装方針であることを示唆します。
ただし、H100のTMAや第4世代Tensor Core、FP8などを最大限使うには、さらに別の最適化余地があります。つまりFlashAttention-2は完成形というより、ハードウェアに合わせてAttention実装を進化させる方向性を示した論文と見るのが自然です。
何に使える？
FlashAttention-2は、Attentionを多用するAIシステム全般で使い道があります。特に、長い系列を扱う開発では効果が出やすい技術です。
長文RAGと社内文書検索
RAGでは、検索した複数文書、引用元、会話履歴、ユーザー指示をまとめてLLMへ渡すことがあります。コンテキストが長くなるほどAttention計算が重くなるため、FlashAttention-2のような正確なAttention高速化は、回答品質を落とさずレイテンシを抑える手段になります。
とくに、法務文書、仕様書、議事録、社内ナレッジのように長文をそのまま扱いたいシステムでは、コンテキストを短く削るだけでは情報落ちが起きます。FlashAttention-2は、入力を削る前に実行基盤側で処理可能な範囲を広げる選択肢になります。
エージェントの長い作業履歴
AIエージェントでは、タスク履歴、ツール実行結果、ファイル差分、エラーログなどが長くなりがちです。履歴を毎回要約する方法もありますが、要約には情報損失があります。
FlashAttention-2のようなAttention高速化が使える環境では、より長い作業履歴をモデルに渡しやすくなります。これは、コード修正、調査、データ分析、複数ステップの業務自動化で、エージェントが前後関係を保つ助けになります。
モデル学習とファインチューニング
長文対応モデルを作る場合、コンテキスト長を伸ばすほど学習コストが上がります。FlashAttention-2は、8kコンテキストのGPT形式学習で大きなスループット改善を示しているため、長文ファインチューニングや継続事前学習のコスト削減に役立ちます。
小規模なチームでも、直接カーネルを実装する必要はありません。PyTorchや主要なLLMフレームワークがFlashAttention系を組み込んでいる場合、設定を正しく選ぶだけで恩恵を受けられることがあります。
画像・音声・動画の長系列処理
論文の動機にもあるように、長い系列はテキストだけの問題ではありません。高解像度画像、音声、動画生成・理解でも、Transformerを使うと系列長が増えます。FlashAttention-2の考え方は、マルチモーダルモデルの学習や推論でも重要です。
特に動画や音声では、時間方向の長い依存関係を扱うため、Attentionの効率化はモデルサイズや入力長の選択肢を広げます。
開発や事業へのヒント
この論文から得られる最大のヒントは、AIプロダクトの性能改善をモデル品質だけで見ないことです。
長文機能はモデル選定だけで決まらない
「長いコンテキストに対応したモデルを使う」だけでは、実用的な長文AI機能は作れません。実際には、Attentionカーネル、KVキャッシュ、バッチング、GPUメモリ、フレームワーク設定が速度とコストを決めます。
FlashAttention-2は、同じモデル構造でも実装次第で大きな性能差が出ることを示しています。長文RAGやコード理解ツールを作るなら、モデルAPIの最大コンテキスト長だけでなく、どのAttention実装が使われているか、どの条件で有効になるかを見る価値があります。
小規模プロダクトでも「入力を削る前に実行効率を見る」
RAGや業務支援ツールでは、コスト削減のために入力文書を短く切り詰める設計をしがちです。しかし、入力削減だけに頼ると、必要な根拠が落ちて回答品質が下がることがあります。
FlashAttention-2の考え方からは、まず実行効率を改善し、そのうえで検索・圧縮・要約を組み合わせる発想が得られます。たとえば、長文を扱う社内AIでは、FlashAttention対応モデルを選び、さらに検索で不要部分を削るという順序のほうが、品質とコストのバランスを取りやすい場合があります。
GPU最適化はプロダクト差別化になり得る
同じLLMを使っていても、推論基盤の効率が違えば、応答速度、同時処理数、価格設計が変わります。FlashAttention-2は低レベルな技術ですが、最終的には「長い入力でも速い」「安くファインチューニングできる」「ユーザーの待ち時間が短い」という事業上の価値につながります。
自社でモデルをホストする場合はもちろん、外部APIを使う場合でも、長文入力を多用する機能ではレイテンシと料金が重要です。Attention高速化の進展を追っておくことは、AI機能の設計判断に役立ちます。
限界
FlashAttention-2は強力ですが、すべての環境で同じ効果が出るわけではありません。
まず、GPUカーネルに深く依存する技術なので、対応ハードウェア、CUDA環境、ライブラリバージョン、head dimension、データ型、mask形式によって使える条件が変わります。論文ではA100を中心に評価されており、他のGPUやCPU環境では同じ改善率を期待できるとは限りません。
次に、FlashAttention-2はAttentionの計算結果を変えずに高速化する技術であり、モデルそのものの推論能力を上げるものではありません。入力が長すぎてモデルが情報を使いこなせない場合や、RAGの検索品質が低い場合は、Attentionが速くなっても回答品質は改善しません。
また、実装の難しさもあります。sliced-Q方式、オンラインsoftmax、ワープ分担、共有メモリ削減などは、通常のアプリケーション開発者が簡単に再実装できるものではありません。実務では、信頼できるライブラリやフレームワークに組み込まれた実装を使うのが現実的です。
さらに、長文処理全体のボトルネックはAttentionだけではありません。KVキャッシュ管理、通信、プリフィルとデコードのスケジューリング、メモリ断片化、分散推論なども効いてきます。FlashAttention-2は重要な部品ですが、長文AIシステム全体を設計するには他の最適化と組み合わせる必要があります。
よくある質問
Q. FlashAttention-2はFlashAttentionと何が違うのですか？
A. 初代FlashAttentionは、Attention行列をHBMに書き出さないことでIOを削減する点が中心でした。FlashAttention-2はその考え方を保ちながら、non-matmul FLOPsの削減、1つのhead内での系列方向並列化、sliced-Qによるワープ分担の改善を加えています。つまり、メモリ効率だけでなくGPUの並列実行効率まで踏み込んで改善しています。
Q. FlashAttention-2を使うとモデルの精度は変わりますか？
A. 基本的には変わりません。FlashAttention-2は正確なAttentionを計算する実装技術であり、疎Attentionのように参照先を削る近似手法ではありません。ただし、実際の学習・推論環境ではデータ型、数値誤差、ライブラリ設定の影響はあるため、導入時には既存出力との差分確認が必要です。
Q. RAGアプリでもFlashAttention-2は役立ちますか？
A. 役立つ可能性があります。RAGでは検索結果、引用、会話履歴、指示文が長くなりやすく、プリフィル時のAttention計算が重くなります。FlashAttention-2対応のモデル実行環境を使うと、長い入力を扱うときのレイテンシやGPUコストを下げられる可能性があります。ただし、検索精度やチャンク設計の問題を直接解決する技術ではありません。
Q. アプリ開発者が自分でFlashAttention-2を実装する必要はありますか？
A. 多くの場合、自分で実装する必要はありません。GPUカーネル最適化は専門性が高いため、実務ではPyTorch、Hugging Face、vLLM、独自推論基盤などでFlashAttention系の実装が有効になっているかを確認するのが現実的です。重要なのは、使っている実行環境がどの条件でFlashAttention-2相当のカーネルを選ぶかを理解することです。
Q. FlashAttention-2だけで長文コンテキスト問題は解決しますか？
A. いいえ。FlashAttention-2はAttention計算を高速化しますが、長文を理解できるかどうかはモデルの学習、位置エンコーディング、検索設計、KVキャッシュ管理、プロンプト構成にも依存します。長文AIシステムでは、FlashAttention-2を基盤の高速化として使い、その上でRAG、要約、キャッシュ、スケジューリングを組み合わせる必要があります。
今日の学び
FlashAttention-2は、長い系列を扱うTransformerでAttention計算がGPU資源を十分に使い切れない課題を扱った論文です。
この課題に対して、non-matmul FLOPsの削減、head内の系列方向並列化、sliced-Qによるワープ分担の改善によって、正確なAttentionのまま高速化する方法を示しました。
ここから得られるヒントは、AIプロダクトの速度とコストはモデル選定だけでなく、実行基盤の細部で大きく変わるということです。長文RAG、エージェント、ファインチューニング、マルチモーダル処理を作るなら、Attentionをどう計算しているかまで見ることが、実用的な性能改善につながります。