FlashAttention（フラッシュアテンション）とは？読み方・LLMのAttention計算をGPUメモリ削減で高速化する仕組み・FlashAttention-3との違いを完全解説

FlashAttentionとは

FlashAttention（フラッシュアテンション）とは、Stanford大学のTri Dao氏らが2022年に提唱したAttention計算の高速化アルゴリズムである。Transformerの中核演算であるScaled Dot-Product Attentionを、GPUのHBM（高帯域メモリ）アクセスを最小化する形で再実装し、従来比2〜4倍の高速化と大幅な省メモリを同時に実現した。FlashAttention-2、FlashAttention-3と進化を続け、2026年現在もLLMの長文コンテキスト化（Llama 3で128K、一部実装では1M）を支える基盤技術となっている。

イメージとしては「Attention計算を、HBM（遅い大容量メモリ）ではなくSRAM（速い小容量メモリ）の中で完結させる工夫」と捉えると分かりやすい。従来のAttentionはQ・K・VとそれらをかけたN×Nの中間行列をHBMに何度も読み書きしていたが、FlashAttentionは行列をブロック分割してSRAMに収め、I/Oを大幅に削減する。ここが重要なポイントです — 数学的には同じ結果を返すが、実行時間とメモリ使用量が劇的に改善されている。

FlashAttentionの読み方

フラッシュアテンション

FlashAttention-3はそのまま「フラッシュアテンションスリー」

FlashAttentionの仕組み

FlashAttentionの中核は、メモリ階層を意識した「Tiling（ブロック分割）」と「Online Softmax（逐次softmax計算）」の2つだ。GPUのSRAMはHBMより圧倒的に高速だが容量が小さいため、Q・K・Vをブロック単位に切り、SRAMの中で部分的なAttention計算を完結させる。

FlashAttentionのブロック分割アプローチ

①Q/K/Vをブロック分割

→

②SRAMにロード

→

③Online softmaxで部分計算

→

④結果を統合してHBMへ書き戻し

Online Softmax

覚えておきたいのは、これは数値計算上のテクニックであり、結果の精度を犠牲にしていないという点だ。

標準のsoftmaxは行ベクトル全体の最大値が必要なため、行列全体をメモリに置く必要があった。Online softmaxは「部分集合だけで暫定的なsoftmaxを計算し、新しいブロックを追加するたびに正規化定数を更新する」という発想で、行列全体をいちどもHBMに置かずに済む。これがFlashAttentionの省メモリの本質だ。

FlashAttention-2 / 3の進化

FlashAttention-2（2023年）は並列化軸の改善でA100の理論ピークに迫る性能を達成。FlashAttention-3（2024年）はNVIDIA H100に最適化され、Tensor CoreとTMA（Tensor Memory Accelerator）の非同期実行を活用し、最大75%のFLOPS利用率（従来は35%）を達成した。FP8への対応も進み、Hadamard変換を回転埋め込みと融合する高度な最適化も加わっている。

FlashAttentionの使い方・実例

基本的な使い方（Quick Start）

# pip install flash-attn
import torch
from flash_attn import flash_attn_func

# Q, K, V を [batch, seq, heads, dim] 形式で用意
q = torch.randn(2, 1024, 16, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 1024, 16, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 1024, 16, 64, device='cuda', dtype=torch.float16)

out = flash_attn_func(q, k, v, causal=True)  # GPT系の因果Attention

主要なフレームワーク（Hugging Face Transformers, vLLM, PyTorch本体のSDPA）はFlashAttentionをデフォルトまたはオプトインで内部利用する設計になっており、明示的に呼び出さなくても恩恵を受けられる。

よくある実装パターン

パターンA: Hugging Face Transformersで有効化

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # ここがポイント
    device_map="auto",
)

向いているケース: 既存のTransformersコードベースを最小変更で高速化したい場合。

避けるべきケース: V100以前の古いGPU（FlashAttentionはAmpere以降が前提）。

パターンB: PyTorch SDPAから自動選択

import torch.nn.functional as F

# PyTorch 2.x はFlashAttentionを自動選択
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

向いているケース: PyTorch標準APIだけで完結させたい開発。コード変更なしでFlashAttentionが選ばれる。

避けるべきケース: マスクや特殊なAttentionパターンが必要で、フォールバック実装に落ちる場合。

アンチパターン: FP32での無理な利用

# ⛔ 推奨されない
out = flash_attn_func(q.float(), k.float(), v.float())

FlashAttentionはFP16/BF16を前提に最適化されており、FP32では理論性能が発揮できない。学習・推論ともにmixed precisionで使うのが正解。FP8対応はFlashAttention-3以降だがH100クラスのハードウェアが前提となる。

FlashAttentionのメリット・デメリット

実務では「導入コストの軽さ」と「対応GPUの世代」を秤にかけて判断するのがセオリーだ。

メリット

標準Attention比で2〜4倍の高速化（FA3はH100でさらに2倍）
メモリ使用量が線形オーダーに削減され、長文コンテキスト学習を実用化
同じ数学結果を返すため、ファインチューニング・推論の精度を損なわない
Hugging Face・vLLM・PyTorch等の主要スタックに統合済み
BSD-3ライセンスでオープン

デメリット

NVIDIA Ampere以降のGPUが前提（V100以前は非対応）
FlashAttention-3の最大効果はH100クラスでないと得られない
カスタムAttentionパターン（block-sparseなど）はサポート範囲外の場合あり
FP32での効果は限定的、低精度（FP16/BF16/FP8）と組み合わせる前提

FlashAttentionとPagedAttention・xFormersの違い

「Attentionを速くする技術」は複数あり、FlashAttention・PagedAttention・xFormersはそれぞれ異なる目的で開発されている。下記の比較表で違いを整理する。

観点	FlashAttention	PagedAttention	xFormers
主な目的	Attention計算の高速化と省メモリ	KVキャッシュのメモリ管理	Attention実装の総合パッケージ
最適化対象	forward/backward両方	推論時のKVキャッシュ	複数のAttention派生実装
学習対応	あり	推論専用	あり
対応ハードウェア	NVIDIA Ampere以降（FA3はHopper）	vLLMが対応する全GPU	NVIDIA中心
代表的な利用先	Hugging Face / PyTorch SDPA	vLLM	Diffusionモデル各種
関係性	基盤技術	FlashAttention上で構築	FlashAttentionを内部利用

つまりFlashAttentionは「Attention計算そのものの高速化」、PagedAttentionは「サービング時のメモリ管理」と階層が異なる。実際にはvLLMがPagedAttentionとFlashAttentionを併用するように、これらは競合ではなく補完関係にある。

よくある誤解

誤解1: 「FlashAttentionを使うと精度が落ちる」

なぜそう誤解されるのか: 「最適化＝近似」という直感が強く、量子化や蒸留と同じ系統の手法と混同されやすい。論文タイトルに「Fast and Accurate」と書かれていることが逆に「速さと引き換えに精度を犠牲にしているのでは」という疑念を生んでいる。

正しい理解: FlashAttentionは数学的には標準Attentionと同じ結果を返す厳密な実装。違いは「どの順序でメモリにアクセスするか」だけで、計算自体は近似していない。Online softmaxも数値計算上の精度差はFP16の丸め誤差レベルに収まる。

誤解2: 「FlashAttentionは推論のみ高速化される」

なぜそう誤解されるのか: vLLMやTGIがFlashAttentionを使って推論を速くする話が広まっている影響で、推論専用の技術と思われがち。背景には「学習はメモリよりFLOPS律速」という思い込みがある。

正しい理解: FlashAttentionは学習（forward+backward）と推論の両方に対応する。長文コンテキストの学習でとくに効果が大きく、Llama 3の128Kコンテキスト学習はFlashAttentionなしでは現実的でなかった。

誤解3: 「どんなGPUでも同じ効果が得られる」

なぜそう誤解されるのか: 「FlashAttentionは標準のAttentionより常に速い」という単純化された情報がインフォグラフィックなどで広まっている。GPUアーキテクチャの違いを踏まえずに語られることが多い。

正しい理解: FlashAttention-2はNVIDIA Ampere（A100）以降を前提、FlashAttention-3はHopper（H100）以降に最適化されている。古いV100やコンシューマGPUでは効果が限定的、もしくは未対応となる場合がある。導入前に対象GPUの世代を必ず確認すること。