FlashAttentionとは?仕組みとFA2・3・4の違い・GPU世代別の実装判断を解説
FlashAttention(フラッシュアテンション)は、Transformerのアテンション計算を、数式を変えずにGPUのメモリ往復だけ減らして高速化するカーネル実装です。2022年5月のarXiv 2205.14135でTri Daoらが発表し、いまはPyTorchとHugging Face Transformersの標準的な選択肢に入っています。この記事では、HBMとSRAMの帯域差という前提、タイリングとオンラインsoftmaxという2つの中核技術、FA1からFA4までの公称スループットと対応GPU、transformers・PyTorch SDPA・vLLMで実際に触る設定キー、そして「入れたのに速くならない」ときに疑う条件までを実装者の目線で整理します。
まとめ:FlashAttentionの仕組みと版選択の結論
FlashAttentionは近似ではありません。標準アテンションと同じ値を返しながら、L×Lのスコア行列をGPUのHBMへ書き出さず、SRAM上のブロック単位で処理しきることで高速化します。論文の理論解析では、HBMアクセス量が標準実装のΘ(Nd+N²)からΘ(N²d²M⁻¹)へ下がります(Nは系列長、dはヘッド次元、MはSRAM容量)。
版の選択はGPU世代でほぼ決まります。Ampere・Ada(A100、RTX 4090)ならFlashAttention-2、Hopper(H100)ならFA3、Blackwell(B200)ならFA4が対応範囲です。PyPIのflash-attnは2.8.3.post1(2026年6月11日)が最新の安定版で、FA4はflash-attn-4の4.0.0b24(2026年7月29日)がpre-release段階にあります。
採用判断の結論を先に書きます。学習・推論のどちらでも、対応GPUを持ちfp16またはbf16で動かせるなら、入れない理由はほぼありません。逆に、fp32での推論が要件だったり、Volta以前のGPUしか持っていなかったりする環境では、導入しても実行時に標準経路へ落ちるだけです。ビルド環境を整えられない場合は、Hugging Face Kernelsの事前ビルド版を先に試すほうが早く結論が出ます。
FlashAttentionが削るもの|HBM往復とO(L²)メモリの内訳
アテンションが遅い理由は、演算量よりもメモリ往復にあります。ここを分けると、何を削っているのかが具体的に見えます。
スコア行列L×LのHBM往復が生む帯域律速というボトルネックの実体
標準的なアテンションは、QKの内積からL×Lのスコア行列を作り、それをHBMへ書き出し、softmaxのために読み戻し、結果をまた書き出し、Vとの積のために再度読み込みます。系列長が8Kなら6,400万要素、fp16でも約128MBの行列が、この往復を何度も通る。演算器は待たされるだけです。この状態を帯域律速と呼びます。アテンションそのものの数式はTransformerの自己注意の仕組みで扱っているため、ここでは削減対象となる往復だけを押さえます。
SRAM192KBと帯域19TB/sというGPUメモリ階層の実測値
削減の余地は、GPUのメモリ階層の落差から生まれます。arXiv 2205.14135のセクション2.1は、A100について「40〜80GBのHBM、帯域1.5〜2.0TB/s」「108基のSMそれぞれに192KBのオンチップSRAM、帯域は19TB/s程度と推定」と記しました。SRAMはHBMより一桁速く、容量は何桁も小さい。192KBの枠にQ・K・Vのブロックを収めきれるなら、中間結果をHBMへ出す必要がなくなります。GPU側の実行モデルはCUDAの仕組みとGPU並列計算に整理しています。
厳密計算を保ったまま省メモリ化する点でスパース化と分かれる境界線
ここを混同すると技術選定を誤ります。FlashAttentionは全トークン対を計算する。省くのはメモリ往復であって、計算するペアではありません。出力は標準アテンションと数値的に一致します(浮動小数点の丸め差を除く)。対してスパースアテンションによる疎化は、見る相手を減らす近似で、精度と引き換えに計算量そのものを落とす手法です。両者は排他ではなく、疎パターンをFlashAttention系カーネルの上に実装する構成も取られます。
タイリングとオンラインsoftmaxによる厳密計算の維持と高速化
中核は2つです。行列をSRAMに収まる大きさへ割るタイリングと、分割しても正しいsoftmaxを出すための逐次計算。順に見ます。
QKVをブロック分割してSRAMへ載せるタイリングの実行手順と制約
タイリングは次の流れで進みます。
- K・Vを行方向のブロックに割り、外側ループで1ブロックずつSRAMへ読み込む
- 内側ループでQのブロックを読み込み、そのブロック対だけのスコアを計算する
- スコアに部分的なsoftmaxを掛け、Vブロックとの積を出力バッファへ足し込む
- 全ブロックを回り終えた時点で、出力バッファが完全なアテンション出力になる
ブロックサイズはSRAM容量とヘッド次元から決まります。制約として効いてくるのがヘッド次元で、公式READMEはFlashAttention-2のサポート範囲を「256まで」としています。ヘッド次元が大きいモデルほどブロックを小さく取らざるを得ず、削減効果も薄まる関係です。
最大値と正規化項を持ち回るオンラインsoftmaxの再スケール処理
softmaxは全要素の最大値と総和を必要とするため、素直に分割すると計算できません。FlashAttentionはブロックごとに「これまでの最大値」と「これまでの正規化項の総和」という2つの統計量を持ち回ります。新しいブロックでより大きな値が現れたら、既に足し込んだ出力を新しい最大値基準へ再スケールしてから加算する。この補正が入るため、分割しても最終的な値は一括計算と一致します。オンラインsoftmaxと呼ばれるこの手法が、厳密性と省メモリを両立させている部分です。
逆伝播でL×L行列を保存せず再計算するメモリ削減のトレードオフ
学習時の勾配計算には、順伝播で作ったスコア行列が要ります。素直に保存すればメモリはO(N²)へ戻ってしまう。FlashAttentionは保存せず、逆伝播の中で統計量からスコアを作り直します。演算量は増える。それでもHBM往復が減るぶん実測は速くなる、という判断です。メモリ使用量は系列長に対して線形になり、論文はPath-X(系列長16K)で正解率61.4%、Path-256(同64K)で63.1%という、従来のTransformerが到達できなかった長さでの学習を報告しました。
FA1からFA4までの世代差|対応GPUと公称スループットの比較
版が上がるたびにアルゴリズムが変わったわけではありません。変わったのは、その世代のGPUが持つ非同期機構をどこまで使い切るかです。
FA1の3倍・FA2の225TFLOPs・FA3の740TFLOPsという公称値
初代(arXiv 2205.14135・2022年5月)は、BERT-largeの学習でMLPerf 1.1の記録比15%短縮、GPT-2で3倍、Long-Range Arenaで2.4倍という数値を示しました。FlashAttention-2(arXiv 2307.08691・2023年7月)は並列化とワーク分割を見直し、初代比で約2倍、A100あたり225TFLOPs(モデルFLOPs利用率72%)に到達。FlashAttention-3(arXiv 2407.08608・2024年7月)はHopperの非同期性とワープ特化を使い、FP16でFA2比1.5〜2.0倍・740TFLOPs(利用率75%)、FP8では1.2PFLOPs近くに届きました。FP8時の数値誤差はベースラインの2.6分の1です。
FA4がBlackwellで到達した1605TFLOPsと2026年3月時点の位置づけ
FlashAttention-4は2026年3月5日に公開されました(arXiv 2603.05451)。B200上のBF16で最大1605TFLOPs(利用率71%)、順伝播はcuDNN 9.13比で1.1〜1.3倍、Triton比で2.1〜2.7倍という数値です。実装はCUTLASSのPython製カーネルDSLであるCuTe-DSLで書き直され、指数関数をSFU経由ではなくFMA側の多項式近似で回す設計に変わりました。ただしPyPIのflash-attn-4は4.0.0b24(2026年7月29日)とpre-release段階にあり、決定的な逆伝播モードは非決定的モードの85〜90%程度のスループットに留まります。本番投入はこの版数表記が取れてから判断する段階です。
AmpereからBlackwellまでのGPU世代と対応版の組み合わせ一覧
手元のGPUで選べる版は、次の対応で確定します。
| GPU世代 | 代表機種 | 選べる版 | パッケージ | データ型 |
|---|---|---|---|---|
| Volta以前 | V100、T4 | 対応なし | 導入不可 | 標準経路のみ |
| Ampere | A100、RTX 3090 | FA2 | flash-attn 2.8系 | fp16、bf16 |
| Ada | RTX 4090、L40S | FA2 | flash-attn 2.8系 | fp16、bf16 |
| Hopper | H100、H800 | FA2、FA3、FA4 | hopper版、4系 | FP8前方向も可 |
| Blackwell | B200 | FA4 | flash-attn-4 | BF16中心 |
真っ先に見るのは1行目です。V100やT4しかない環境では、どのパッケージを入れても実行時に標準経路へ落ちます。AMD側はROCm 6.0以上でMI200・MI300系がfp16・bf16に対応しました。
実装で触る場所|transformers・PyTorch SDPA・vLLMの設定
論文の図よりも、実務で詰まるのは設定キーとカーネルの対応関係です。触る箇所は学習側と推論側で分かれます。
transformersのattn_implementationで指定する4つの値
Hugging Face Transformers(v5.14.0・2026年8月時点)では、アテンション実装はattn_implementationという1つの引数に集約されています。指定できる主な値はflash_attention_2、flash_attention_3、sdpa、flex_attentionの4つ。それぞれにpagedを前置した派生も用意されています。from_pretrainedで渡すのが基本で、モデルを読み直さずに切り替えたい場合はset_attn_implementationを呼びます。マルチモーダルモデルでは辞書を渡し、視覚側だけsdpaに落とす構成も公式ドキュメントに示されました。
PyTorchのsdpa_kernelでFLASH_ATTENTIONを強制する手順
PyTorchのscaled_dot_product_attentionは、CUDA上では条件を見て内部実装を自動で選びます。固定したいときはtorch.nn.attentionのsdpa_kernelコンテキストマネージャでSDPBackend.FLASH_ATTENTIONを明示する。条件を満たさないまま強制すると例外が上がるため、黙って遅い経路へ落ちる状態を切り分ける手段になります。ベンチマークの前後でこれを挟み、想定した経路が本当に通っているかを先に確定させてください。
vLLMのVLLM_ATTENTION_BACKENDで指定する推論側の設定
推論サーバ側では、バックエンド名で選択します。vLLMは環境変数VLLM_ATTENTION_BACKENDを持ち、FLASH_ATTNのほかTORCH_SDPA、XFORMERS、FLASHINFER、FLASHMLAなどを受け付けます。未指定なら対応状況を見て自動選択されるため、通常は触りません。触るのは、FP8のKVキャッシュを使うなど機能の組み合わせで自動選択が意図とずれるときです。サーバ全体の構成と、PagedAttentionによるKVキャッシュ管理との役割分担はvLLMの仕組みと使い方で解説しています。
kernels-communityの事前ビルド版でCUDA不一致を避ける手順
導入で最も時間を溶かすのがビルドです。pip install flash-attn --no-build-isolationはCUDAツールキットとninjaを要求し、環境によっては1時間近くかかります。回避経路がHugging Face Kernelsによる事前ビルド済みカーネルの取得で、attn_implementationにkernels-community配下のflash-attn2を指定すると、Hubから実行時にカーネルが降ってきてAttentionInterfaceへ自動登録される。flash-attnパッケージそのものを入れる必要はありません。
導入しても速くならない原因|head_dim・dtype・マスク形式
「入れたのに変わらない」という相談は、ほぼ3つの原因に収束します。いずれも設定は通っているのに、実行時にカーネルが選ばれていません。
head_dim256超とfp32入力でカーネルが黙って外れる条件
FlashAttention-2が扱えるデータ型はfp16とbf16だけで、fp32は対象外です。bf16はAmpere以降に限られます。ヘッド次元は256までという上限もあり、これを超えるモデルでは選択されません。厄介なのは、多くのフレームワークがこの不一致を例外ではなく静かなフォールバックとして処理する点です。ログには何も出ないまま標準経路で走り、ベンチマークだけが「効果なし」を示します。まずmodel.dtypeとconfig.head_dimを出力して、前提条件を数値で確認してください。
4Dマスクとpaddingの扱いがFA系だけ異なるという実装上の差
マスクの形式もフォールバックの引き金になります。Transformersの公式ドキュメントは、sdpaがboolまたはfloatの4Dマスクを受け取るのに対し、flash_attention_2とflex_attentionは独自形式(前者は2Dのパディングマスク、後者はBlockMask)しか受け付けないと明記しました。自前で4Dマスクを組んで渡している既存コードをFA2へ切り替えると、ここで噛み合いません。系列長の分散が大きいワークロードでは、パディングを畳む処理と組み合わせて初めて効果が出ます。
短系列とバッチ小でFlashAttentionの効果が出ない実測条件
前提条件を満たしていても、系列が短ければ差は出ません。削減対象がL×L行列のHBM往復である以上、Lが小さいほど削る余地も小さくなる。系列長512程度の分類タスクでは、測定誤差に埋もれる水準に留まります。デコード段階も同様です。1トークンずつ生成する局面ではQ側の長さが1で、律速はKVキャッシュの読み出しとメモリ削減へ移り、FlashAttentionの守備範囲から外れます。効果が明確に出るのは、系列長2K以上のプリフィルと学習です。
FlashAttentionを採用してよい条件と、見送るべき場面
ここは言い切ります。FlashAttentionは「長文を扱うなら検討する」類の改造ではなく、条件を満たすなら既定で入れておく実装です。近似ではないため、精度側のリスクを取らずに済みます。
採用してよい条件=GPU世代・系列長・dtypeの3点の実測確認
第1に、GPUがAmpere以降であること。A100、RTX 3090、RTX 4090、L40S、H100のいずれかなら該当します。第2に、学習またはプリフィルで扱う系列長の中央値が2Kを超えていること。平均ではなく中央値で見る。第3に、モデルをfp16またはbf16で動かせること。混合精度の学習では既に満たしている場合がほとんどです。3点が揃うならattn_implementationの切り替えだけで済み、空いたメモリでバッチサイズを増やせばスループット側にも二重に効きます。
見送るべき場面=fp32推論・Volta以前・ビルド不可の環境
見送る場面は3つです。第1に、fp32での推論が要件に入っている場合。数値の完全再現が求められる検証系や、fp32前提の画像バックボーンがこれに当たります。第2に、V100やT4しか調達できない環境。パッケージだけ入れても標準経路へ落ち、依存関係の管理コストが増えるだけです。第3に、CUDAツールキットを持ち込めずビルドが通らない場合。ただし3つ目は前述の事前ビルド版で解消する余地があり、切り捨てる前に一度試す価値があります。
LLM推論基盤の構築を外注するとき見積書で確認する技術要件5項目
アテンションカーネルまで含めた学習・推論基盤を外注する場合、見積書で確認すべき項目は次の5つです。
- 採用するアテンション実装と版が明記され、調達GPUの世代と対応づいているか
- fp16・bf16・FP8のどれで動かす前提か、精度検証の工程が含まれているか
- カーネルの選択結果をログまたはベンチマークで確認する手順が入っているか
- 系列長の実測分布に基づいて、プリフィルとデコードの費用が分けられているか
- ビルド環境(CUDA版・ツールキット)の前提が納品環境と一致しているか
この5項目が埋まらない見積書は、カーネルが実行時に外れていても検収を通ってしまう。GPU調達の段階から判断が要る場合は、機械学習モデル開発の受託で系列長の実測と構成の選定から相談できます。導入の全体像はLLMの仕組みと企業導入の判断基準で先に押さえておくと前提が揃います。
よくある質問
FlashAttentionの導入検討で実際に問い合わせの多い論点を5つ挙げます。
FlashAttentionを使うと出力の精度は落ちますか?
落ちません。FlashAttentionは全トークン対を計算する厳密な実装で、標準アテンションと同じ値を返します。異なるのは計算の順序とメモリの使い方だけです。浮動小数点の加算順序が変わるため最終ビットまで一致するとは限りませんが、これは精度低下とは別の話です。精度と引き換えに計算量を落とすのはスパースアテンションなどの近似手法で、FlashAttentionはその系統には属しません。
FlashAttention-2と3のどちらを入れればよいですか?
GPUで決まります。A100やRTX 4090などAmpere・Ada世代ではFA3が動かないため、選択肢はFA2だけです。H100・H800ならFA3が使え、FP16でFA2比1.5〜2.0倍という公称値が出ています。ただしFA3はリポジトリ上ベータ扱いで、CUDA 12.3以上(12.8推奨)を要求し、FP8は前方向のみの対応。学習で逆伝播までFP8を通したい要件があるなら、現時点ではFA2に留めるほうが安全に運べます。
FlashAttentionとPagedAttentionは何が違うのですか?
削減する対象が違います。FlashAttentionはアテンション計算そのもののメモリ往復を減らすカーネルで、主にプリフィルと学習に効きます。PagedAttentionはvLLMが導入したKVキャッシュの管理方式で、キャッシュをページ単位に分けて断片化を抑える仕組みです。効くのはデコード段階のメモリ効率とバッチ密度になります。両者は競合せず、vLLMはFLASH_ATTNバックエンドとPagedAttentionを同時に使えます。詳細はvLLMとPagedAttentionの解説を参照してください。
flash-attnのインストールに時間がかかるのはなぜですか?
CUDAカーネルをソースからコンパイルするためです。pip install flash-attn --no-build-isolationはCUDAツールキットとninjaを必要とし、並列度が足りないと数十分から1時間規模になる。回避策は2つあります。Hugging Face Kernelsの事前ビルド版をHubから取得する経路と、PyTorchのSDPA経由で内蔵実装を使う経路です。
推論の遅延を下げたいのですが、まず何から測ればよいですか?
プリフィルとデコードの時間配分からです。入力が長くプリフィルが支配的ならFlashAttentionが効く一方、出力が長くデコードが支配的なら効果は限定的で、量子化やバッチ設計のほうが効きます。実測せずにカーネルだけ差し替えても、改善幅は予想を外す。推論全体でどこを詰めるかの整理はLLM推論の高速化手法と推論基盤の選び方にまとめています。
関連記事
- Transformerとは?自己注意の仕組みとBERT・GPTの違いを実装目線で解説:FlashAttentionが高速化する対象である自己注意の数式を扱っています
- スパースアテンションとは?疎なAttentionの仕組みとNSA・DSAの実装判断を解説:厳密計算のFlashAttentionと対をなす近似側の手法です
- KVキャッシュとは?LLM推論を高速化する仕組みとメモリ削減・実装判断を実装者向けに解説【2026年版】:デコード段階の律速要因はこちらで扱っています
- vLLMとは?PagedAttentionの仕組み・使い方とOllama・TensorRT-LLMとの違いを実装者目線で解説:推論サーバ側でバックエンドを選ぶ全体像が分かります
- LLM推論とは?仕組みと高速化の手法・推論基盤の選び方を実装視点で解説【2026年】:カーネル以外の高速化手段との優先順位を整理しています