FlashAttentionの仕組み｜タイル化とオンラインsoftmaxを自作

この記事で実際に動かしたもの：numpy だけで書いた素朴な注意・FlashAttention（論文の Algorithm 1 の順序）・FlashAttention-2 の順序の注意・因果マスク・逆伝播（再計算）と、HBM の読み書きの量を数える模型。GPU では動かしていません。GPU の速さの数字はすべて原論文からの引用で、手元の numpy の時間とは別のものです。

Transformer の注意は、系列長 n のとき n×n のスコア行列を作ります。n が2倍になると行列は4倍になり、長い文脈ではこれがメモリと時間の重さになります。FlashAttention はこの行列をメモリに書き出さずに、小さなブロックを高速なオンチップのメモリ（SRAM）に載せたまま softmax と掛け算まで済ませます。鍵になるのは「softmax は全部の値を見ないと計算できない」という思い込みを外す、オンライン softmax の式です。

flowchart TD
    A[Q を Br 行ずつに分ける] --> B[行ブロック Qi を載せる]
    B --> C[K と V を Bc 行ずつ読む]
    C --> D[Sij = Qi と Kj の内積]
    D --> E[最大値 m と和 l を更新<br>途中の出力を補正して足す]
    E -->|次の Kj| C
    E -->|最後の Kj まで| F[Oi = 途中の出力 / l<br>L = m + log l を保存]
    F -->|次の Qi| B

よく調べられることの答えを先にまとめます。

確認環境：Ubuntu 24.04.5 LTS（Intel Xeon 2.10GHz・2コア・L2 4MiB、ほかの作業と同時に動いている共有のマシン）で、Python 3.13.18 と numpy 2.5.3（OpenBLAS 0.3.34、OPENBLAS_NUM_THREADS=1）だけを使いました。PyTorch・CUDA・GPU は使っていません。Q・K・V は標準正規分布の乱数（ヘッド次元 d=64 が中心）で、学習済みのモデルの重みは使っていません。メモリは Python 標準の tracemalloc で numpy の配列の確保量のピークを測り、時間は3〜5回の中央値です（共有のマシンのため10〜20％ぶれます）。FlashAttention・FlashAttention-2・FlashAttention-3 の主張と数字は arXiv の原論文、オンライン softmax の先行研究は Milakov・Gimelshein（2018）と Rabe・Staats（2021）の原論文、使い方は flash-attn の GitHub の README・PyTorch 2.15 の公式ドキュメント・Hugging Face Transformers の公式ドキュメントで確かめました。図は mermaid で描画を確かめています（確認日 2026年11月4日）。

FlashAttentionとは（何を変えて、何を変えないか）

この節では、FlashAttention が注意の計算のどこを変え、どこを変えないのかが分かります。

FlashAttention は、Tri Dao ほかが2022年に発表した「IO を意識した厳密な注意のアルゴリズム」です（arXiv:2205.14135）。論文の題名は「FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness」で、Exact（厳密）とあるとおり、計算する値は softmax(QKᵀ/√d)V そのものです。注意を間引く近似（スパース化・低ランク化）とは別の種類の工夫です。

変えないものは、計算量（FLOPs）と結果です。論文の定理1は、FlashAttention が O(N²d) の FLOPs で、入力と出力のほかに O(N) の追加メモリしか使わないことを示しています（N は系列長、d はヘッド次元）。速くなる理由は計算を減らしたからではなく、遅いメモリへの読み書きを減らしたからです。この点は後の「numpy では速くならない理由」の節で、読み書きの量を数えて確かめます。

素朴な注意はなぜメモリを食うのか（n×n の行列の実測）

この節では、素朴な注意が系列長 n に対して n² のメモリを使う様子と、numpy で書くときの dtype の落とし穴が分かります。

まず比べる相手として、素朴な注意を numpy で書きます。1ヘッド・バッチ1として、Q・K・V はどれも n×d の行列です。

import numpy as np

def naive_attention(Q, K, V, causal=False):
    """素朴な注意：n×n の S と P を丸ごと作る"""
    scale = Q.dtype.type(1.0 / np.sqrt(Q.shape[-1]))  # dtype をそろえる（後述）
    S = (Q @ K.T) * scale                      # n×n
    if causal:
        n = S.shape[0]
        S = np.where(np.tril(np.ones((n, n), dtype=bool)), S, -np.inf)
    m = S.max(axis=1, keepdims=True)
    P = np.exp(S - m)                          # n×n
    P /= P.sum(axis=1, keepdims=True)
    return P @ V

この関数を float32・d=64 で、tracemalloc を使って numpy の配列の確保量のピークを測りました。「np.sqrt(d) で割る書き方」は、(Q @ K.T) / np.sqrt(d) とだけ書いたものです。

import tracemalloc
def peak_mib(f, *args):
    tracemalloc.start(); tracemalloc.reset_peak()
    base = tracemalloc.get_traced_memory()[0]
    out = f(*args)
    peak = tracemalloc.get_traced_memory()[1] - base
    tracemalloc.stop()
    return peak / 2**20

どの書き方でも、n が2倍になるとメモリは4倍になります。in-place で一時配列を減らしても n×n の1枚は消えません。n=16,384 なら float32 の1枚で1GiB、ヘッド数とバッチ数の分だけさらに増えます。

表の左の列が真ん中の列のちょうど2倍なのは、numpy 2 の型の規則のためです。np.sqrt(64) は Python の float ではなく numpy.float64 を返し、numpy 2（NEP 50）では float32 の配列を numpy.float64 のスカラーで割ると結果が float64 になります。手元で確かめると次のようになりました。

type(np.sqrt(64))                 -> <class 'numpy.float64'>
(Q32 @ K32.T / np.sqrt(64)).dtype -> float64
(Q32 @ K32.T / 8.0).dtype         -> float32
n=4096 の時間: np.sqrt(d) で割る 1114ms / dtype をそろえる 554ms / in-place 385ms

エラーも警告も出ないまま、メモリも時間も約2倍になります。この記事のコードは Q.dtype.type(1.0 / np.sqrt(d)) で scale の型を入力にそろえています。

オンライン softmax とは（式の導出）

この節では、softmax を値を1つずつ（またはブロックずつ）見ながら1パスで計算できる理由と、その式が分かります。

オンライン softmax は、最大値 m と指数の和 l を逐次更新し、最大値が変わったときにそれまでの和を exp(m旧−m新) 倍して合わせることで、1回の走査で正しい正規化の項を得る方法です。Milakov・Gimelshein「Online normalizer calculation for softmax」（arXiv:1805.02867、2018）が示し、FlashAttention の論文も付録 B.1 でこの論文を引いています。

1. なぜ最大値を引くのか

softmax は softmax(x)_k = exp(x_k) / Σ exp(x_j) です。そのまま計算すると、float32 では exp(88.8) が inf になります。どの x_j からも同じ数 m を引いても、分子と分母に同じ exp(−m) が掛かるだけなので値は変わりません。ふつうは m = max(x) を引いて、exp の引数を0以下にします。

exp そのまま: [nan nan nan]          (x = [100, 101, 102], float32)
RuntimeWarning: overflow encountered in exp
RuntimeWarning: invalid value encountered in divide
最大値を引く: [0.09003057 0.24472846 0.66524094]
float32: exp(88.7) = 3.325977e+38, exp(88.8) = inf
float16: exp(11.0) = 5.987e+04,   exp(11.1) = inf

問題は、m = max(x) を知るには全部の値を一度見る必要があることです。素朴に書くと「最大値を求める→指数の和を求める→割る」の3回の走査になり、注意ではそのたびに n×n の行列を読むことになります。

2. 1つずつ更新する式

k 番目までの最大値を m_k、k 番目までの和を l_k = Σ_{j≤k} exp(x_j − m_k) と定義します。新しい値 x_k が来たときは次のように更新します。

m_k = max(m_{k-1}, x_k)
l_k = l_{k-1} * exp(m_{k-1} - m_k) + exp(x_k - m_k)

l_{k−1} は m_{k−1} を引いた指数の和なので、exp(m_{k−1} − m_k) を掛けると「m_k を引いた指数の和」に直ります。そこに新しい項を足せば、定義どおり l_k になります（k についての帰納法）。最大値が変わらなければ補正の係数は exp(0)=1 です。乱数1万個で確かめた結果が次です。

(k, x_k, m_k, l_k): (0, 1.728, 1.728, 1.0) (1, 4.108, 4.108, 1.09253)
                    (2, 1.652, 4.108, 1.17832) (3, -6.516, 4.108, 1.17834)
                    (4, 4.527, 4.527, 1.77524)
online l = 2.517961843168063   2パスの和 = 2.517961843168072   相対誤差 3.5e-15

k=1 で最大値が 1.728 から 4.108 に変わったとき、それまでの和 1.0 は exp(1.728−4.108)≈0.0925 倍されてから、新しい項 1.0 が足されています。

3. 2つのブロックを合わせる式（注意への拡張）

注意では、softmax の重みで V の行を足し合わせた出力も必要です。行を2つのブロックに分け、それぞれで「最大値 m_a・和 l_a・正規化していない出力 acc_a = Σ exp(s_j − m_a) v_j」を求めておくと、全体は次の式で合わせられます。

m   = max(m_a, m_b)
l   = exp(m_a - m) * l_a   + exp(m_b - m) * l_b
acc = exp(m_a - m) * acc_a + exp(m_b - m) * acc_b
o   = acc / l

FlashAttention の論文の3.1節は、この形を2ブロックの例で示し、「追加の統計 m(x)・ℓ(x) を持っておけば softmax を1ブロックずつ計算できる」と述べています。この結合は順番を入れ替えても、木の形に組んでも同じ答えになります。1万個を10個のブロックに分けて試しました。

順番 [0,1,2,3,4,5,6,7,8,9]   max|err| 7.77e-16
順番 [1,9,7,0,8,4,3,6,5,2]   max|err| 5.55e-16
木の形で結合               max|err| 4.44e-16

順番に依らないことは、GPU で K・V の方向にも仕事を分け、最後に部分結果を合わせる実装（推論時の長い文脈など）の根拠にもなります。

タイル化した注意を numpy で書く（FlashAttention の Algorithm 1）

この節では、FlashAttention の論文の Algorithm 1 と同じ順序のループを numpy で書いたものと、各行の意味が分かります。

論文の Algorithm 1 は、外側のループで K・V のブロック（Bc 行）を、内側のループで Q のブロック（Br 行）を回します。ブロックの大きさは SRAM の大きさ M から Bc = ⌈M/4d⌉、Br = min(⌈M/4d⌉, d) と決めます。出力 O と統計 m・l は HBM 上にあり、内側のループのたびに読み書きします。

def flash_attention_v1(Q, K, V, Br=64, Bc=64):
    """FlashAttention（論文の Algorithm 1）：外側が K,V のブロック。毎回 O を正規化し直す"""
    n, d = Q.shape
    scale = Q.dtype.type(1.0 / np.sqrt(d))
    O = np.zeros((n, d), dtype=Q.dtype)
    l = np.zeros((n, 1), dtype=Q.dtype)
    m = np.full((n, 1), -np.inf, dtype=Q.dtype)
    for j in range(0, n, Bc):
        Kj, Vj = K[j:j+Bc], V[j:j+Bc]
        for i in range(0, n, Br):
            Sij = (Q[i:i+Br] @ Kj.T) * scale
            mij = Sij.max(axis=1, keepdims=True)
            Pij = np.exp(Sij - mij)
            lij = Pij.sum(axis=1, keepdims=True)
            mi, li = m[i:i+Br], l[i:i+Br]
            mnew = np.maximum(mi, mij)
            a, b = np.exp(mi - mnew), np.exp(mij - mnew)
            lnew = a * li + b * lij
            O[i:i+Br] = (li * a * O[i:i+Br] + b * (Pij @ Vj)) / lnew
            m[i:i+Br], l[i:i+Br] = mnew, lnew
    return O

この順序では、O を正規化済みの形で持つため、ブロックごとに掛け算と割り算が増えます。FlashAttention-2 はここを減らしました。

FlashAttention-2 の順序で書き直す（外側を Q のブロックに）

この節では、FlashAttention-2 の論文が変えたループの順序と、正規化を最後に1回だけ行う書き方が分かります。

FlashAttention-2（Tri Dao、arXiv:2307.08691）の3.1節と3.2節は、次の2つを変えています。1つ目は、出力を正規化していない形のまま持ち、ループの最後に1回だけ diag(ℓ)⁻¹ で割ること。2つ目は、外側のループを Q の行ブロック、内側を K・V の列ブロックにしたことです（論文は、この考えを Phil Tillet の Triton の実装によるものとしています）。さらに逆伝播用には m と ℓ の両方ではなく、logsumexp L = m + log(ℓ) だけを保存します。

def flash_attention_v2(Q, K, V, Br=64, Bc=64, causal=False, return_lse=False):
    """FlashAttention-2 の順序：外側が Q のブロック。O は最後に1回だけ割る。logsumexp を返す"""
    n, d = Q.shape
    scale = Q.dtype.type(1.0 / np.sqrt(d))
    O = np.empty((n, d), dtype=Q.dtype)
    L = np.empty(n, dtype=Q.dtype)
    for i in range(0, n, Br):
        Qi = Q[i:i+Br]
        r = Qi.shape[0]
        acc = np.zeros((r, d), dtype=Q.dtype)
        mi = np.full((r, 1), -np.inf, dtype=Q.dtype)
        li = np.zeros((r, 1), dtype=Q.dtype)
        jend = min(n, i + r) if causal else n
        for j in range(0, jend, Bc):
            Sij = (Qi @ K[j:j+Bc].T) * scale
            if causal and j + Bc > i:          # 対角にかかるブロックだけマスク
                rows = np.arange(i, i + r)[:, None]
                cols = np.arange(j, j + Sij.shape[1])[None, :]
                Sij = np.where(cols <= rows, Sij, -np.inf)
            mnew = np.maximum(mi, Sij.max(axis=1, keepdims=True))
            Pij = np.exp(Sij - mnew)
            alpha = np.exp(mi - mnew)
            li = alpha * li + Pij.sum(axis=1, keepdims=True)
            acc = alpha * acc + Pij @ V[j:j+Bc]
            mi = mnew
        O[i:i+r] = acc / li
        L[i:i+r] = (mi + np.log(li))[:, 0]
    return (O, L) if return_lse else O

外側が Q のブロックなので、行ブロックごとの acc・m・ℓ は内側のループの間ずっと手元（GPU ではレジスタや SRAM）に置いたままにでき、HBM に書くのは最後の O と L だけです。GPU では行ブロックごとに別のスレッドブロックへ割り当てられるので、FlashAttention-2 はバッチとヘッドに加えて系列長の方向にも並列化しています（3.2節）。

出力の一致とメモリの実測（n² から n へ）

この節では、タイル化した注意が素朴な注意と同じ値を出すこと、追加のメモリが系列長に対してほぼ一定になることを実測で確かめられます。

出力の一致

float64 の素朴な注意を基準にして、最大絶対誤差を測りました。ブロックの大きさは Br=Bc=64 です。

float32 では、タイル化しても素朴な計算と同じ水準の丸め誤差で、float64 では機械イプシロンの数倍に収まりました。ブロックの大きさを (Br, Bc) = (16, 128)・(128, 16)・(1000, 1000) に変えても float64 で 5e-16 以下、n=200 で Br=Bc=1（1要素ずつ）にしても 1.1e-15 でした。端数のブロック（n=1,000 を 64 で割ると余りが出る）でも結果は変わりません。

メモリ

同じ tracemalloc の測り方で、float32・d=64 のメモリのピークを比べました。出力 O（n×d）はどの方法でも必要なので、FlashAttention-2 の順序の列は「出力を除いた分」も示します。

出力を除いた分は、ブロックの大きさで決まる一時配列（Br×Bc のスコアと確率、Br×d の acc）と、n 個の logsumexp L だけです。n を512から16,384へ32倍にしても 0.098MiB から 0.159MiB にしか増えず、増えた 0.06MiB はほぼ L の float32 16,384個分（0.0625MiB）です。論文の定理1の「入力と出力のほかに O(N) の追加メモリ」と合います。8,192 以上の素朴な注意は、このマシンの1プロセスの上限（1GB）を超えるおそれがあるため測っていません。

numpy では速くならない理由（HBM の読み書きを数える）

この節では、手元の numpy でタイル化しても速くならない理由と、GPU で速くなる理由を、時間の実測と読み書きの量の数え上げで分けて理解できます。

手元の numpy の時間

float32・d=64 で、ブロックの大きさを変えて時間を測りました（3〜5回の中央値。この数字は CPU 1スレッドの numpy の時間で、GPU の FlashAttention の速さとは関係がありません）。

ブロックが小さいと、numpy の関数呼び出しと Python のループの手間がブロックの数だけかかります。B=16・n=4,096 では1ブロックあたり約63マイクロ秒で、計算そのものより呼び出しの手間が大きくなっています。B=128〜256 まで大きくすると素朴な計算と同じくらいになりますが、それ以上は速くなりません。理由は3つあります。

論文が示すメモリ階層と、読み書きの量

FlashAttention の論文の2.1節は、A100 の HBM を 40〜80GB・帯域 1.5〜2.0TB/s、オンチップの SRAM を108個のストリーミングマルチプロセッサ（SM）それぞれに 192KB・帯域は推定約 19TB/s としています。SRAM は1桁速いが、はるかに小さいという関係です。そのうえで定理2として、標準の注意の HBM アクセスは Θ(Nd + N²)、FlashAttention は Θ(N²d²M⁻¹)（d ≤ M ≤ Nd）であることを示し、命題3で、すべての SRAM の大きさについてこれより漸近的に少ない厳密な方法は無いという下限も示しています。

この式の意味を見るために、論文の Algorithm 0（標準の注意）と Algorithm 1 の読み書きを要素の数で数える模型を書きました。M は SRAM 192KB を fp16 の要素数（98,304）に直し、ブロックの大きさは論文どおり Bc=⌈M/4d⌉、Br=min(Bc, d) です。これは手元の計算での数え上げで、GPU で測った値ではありません。

def std(N, d):   # Algorithm 0：Q,K 読み・S 書き / S 読み・P 書き / P,V 読み・O 書き
    return 2*N*d + N*N + 2*N*N + N*N + N*d + N*d
def fa1(N, d, M):
    Bc = math.ceil(M / (4*d)); Br = min(Bc, d); c = N*d + 2*N   # O,l,m の初期化
    for j in range(0, N, Bc):
        c += 2 * min(Bc, N-j) * d                               # Kj, Vj 読み
        for i in range(0, N, Br):
            br = min(Br, N-i)
            c += 2*br*d + 2*br                                  # Qi, Oi, li, mi 読み
            c += br*d + 2*br                                    # Oi, li, mi 書き
    return c

どちらも N² に比例して増えますが、係数が違います。標準は N² の行列を4回読み書きし、Algorithm 1 は Q・O を Tc=N/Bc 回読み書きするので、およそ 12N²d²/M になります。論文の図2の例では、HBM の読み書きが 40.3GB から 4.4GB に減り、実行時間も短くなった例を示しています（1章では「最大9倍少ない HBM アクセス」）。

同じ模型で d と M を変えると、効き方の条件が見えます（N=4,096）。

d²/M が小さいほど得で、d が大きい・SRAM が小さいと、この単純な模型では標準より読み書きが増えます。実際のカーネルは d ごとにブロックの大きさを調整しており（FlashAttention-2 の3.3節は、ブロックの大きさをヘッド次元ごとに手で調整したと述べています）、この模型の数字がそのまま GPU の速さになるわけではありません。読み取れるのは「FlashAttention の得は、計算を減らすことではなく、SRAM に収まる大きさのブロックで HBM の往復を減らすことから来る」という構造です。

論文が報告する速さ（GPU、引用）

以下はすべて FlashAttention の原論文（arXiv:2205.14135）の報告で、手元では測っていません。

• BERT-large の学習：MLPerf 1.1 の学習速度の記録に対して15％速い（表1。17.4±1.4分と20.0±1.5分）
• GPT-2 の学習：Hugging Face の実装に対して最大3倍（表2。small 3.5倍・medium 3.0倍）、Megatron-LM に対して 1.7〜1.8倍
• Long-Range Arena：標準の Transformer に対して2.4倍（表3）
• 注意そのもの：GPT-2 の条件で PyTorch の実装に対して最大7.6倍（図1右）
• メモリ：厳密な注意の比較対象に対して最大20倍少ない（4.3節）
• 長い文脈：Path-X（系列長16K）で 61.4％、ブロックスパース版で Path-256（系列長64K）で 63.1％（表6）

因果マスクとの組み合わせ（ブロックを飛ばす）

この節では、デコーダの因果マスク（causal）をタイル化した注意に入れる方法と、計算を約半分にできる理由が分かります。

因果マスクでは、位置 i の行は j ≤ i の列しか見ません。ブロックで考えると、全部の列が行より右にあるブロックは丸ごと −∞ なので、計算せずに飛ばせます。対角にかかるブロックだけ要素ごとにマスクし、対角より左のブロックはマスクが要りません。FlashAttention-2 の論文の3.1.1節は、系列長が長いときに飛ばせるブロックがおよそ半分で、因果マスクの無い注意に比べて約1.7〜1.8倍速くなったと報告しています（GPU）。

行ブロック数を T とすると計算するのは T(T+1)/2 個で、T が大きいほど半分に近づきます。手元の numpy でも、n=4,096・B=128 で因果マスクなし 472.1ms に対して因果マスクあり 273.3ms（約1.7倍）でした。一方、素朴な注意に因果マスクを足すと、n×n のマスクを作って np.where を通す分だけ遅くなり、490.0ms が 726.0ms になりました。素朴な計算はマスクした要素も計算してから捨てるためです。出力は float64 で素朴な計算と 7.8e-16 の差で一致し、Br=48・Bc=80 のように行と列のブロックの大きさが違っても同じでした。

行が全部マスクされると NaN になる

パディングのマスクでは、ある行ブロックから見て最初のブロックの列が全部無効になることがあります。このとき行の最大値が −∞ のままで、exp(−∞ − (−∞)) が NaN になります。先頭5トークンをパディングにした n=12・B=4 の例で、対策の有無を比べました。

RuntimeWarning: invalid value encountered in subtract
  P=np.exp(S-ms); a=np.exp(mi-ms)
対策なし NaN の数: 96          (出力 12×8 の全要素)
対策あり NaN の数: 0  max|err| 3.33e-16

最初のブロックで NaN が入ると、その後に有効な列が来ても NaN が残り、出力の全部が NaN になりました。対策は、指数をとるときに使う最大値が −∞ なら0に置き換え、最後に ℓ=0 の行（全部マスクの行）は0で割らないようにすることです。

mn = np.maximum(mi, S.max(axis=1, keepdims=True))
ms = np.where(np.isneginf(mn), 0.0, mn)   # 行が全部マスクのときは 0 を引く
P = np.exp(S - ms); a = np.exp(mi - ms)
li = a * li + P.sum(axis=1, keepdims=True); acc = a * acc + P @ Vj; mi = mn
# 最後：O = acc / np.where(li == 0, 1, li)

逆伝播は保存せずに作り直す（再計算）

この節では、FlashAttention が学習時の逆伝播で n×n の行列を保存しない方法と、その勾配が正しいことの確かめ方が分かります。

素朴な実装は、順伝播で作った確率 P（n×n）を逆伝播のために保存します。FlashAttention の論文の3.1節は、O と softmax の統計だけを保存し、逆伝播で Q・K・V のブロックから S と P を作り直す方法を示し、これを「選択的な勾配チェックポイントの一種」と呼んでいます。FlashAttention-2 では、保存する統計が logsumexp L だけになります。L があれば P = exp(S − L) で正規化済みの確率がそのまま作れるからです。

def flash_backward(Q, K, V, O, L, dO, Br=64, Bc=64):
    """P を保存せず、Q・K と logsumexp L からブロックごとに作り直す"""
    n, d = Q.shape
    scale = Q.dtype.type(1.0 / np.sqrt(d))
    D = (dO * O).sum(1)                         # rowsum(dO∘O)
    dQ = np.zeros_like(Q); dK = np.zeros_like(K); dV = np.zeros_like(V)
    for j in range(0, n, Bc):
        Kj, Vj = K[j:j+Bc], V[j:j+Bc]
        for i in range(0, n, Br):
            P = np.exp(Q[i:i+Br] @ Kj.T * scale - L[i:i+Br, None])   # 再計算
            dV[j:j+Bc] += P.T @ dO[i:i+Br]
            dP = dO[i:i+Br] @ Vj.T
            dS = P * (dP - D[i:i+Br, None])
            dQ[i:i+Br] += dS @ Kj * scale
            dK[j:j+Bc] += dS.T @ Q[i:i+Br] * scale
    return dQ, dK, dV

softmax の逆伝播は dS = P∘(dP − rowsum(dP∘P)) ですが、rowsum(dP∘P) = rowsum(dO∘O) が成り立つので、行ごとの D を最初に O と dO から求めておけば、ブロックの中だけで dS が作れます。n=512・d=64・float64 で、P を保存する素朴な逆伝播と比べ、さらに dQ の3要素を中心差分（幅1e-6）で確かめました。

dQ max|差| 5.55e-16
dK max|差| 4.72e-16
dV max|差| 4.16e-16
dQ[3,5]    有限差分  0.03102810  flash  0.03102810
dQ[100,0]  有限差分 -0.06534607  flash -0.06534607
dQ[511,63] 有限差分 -0.03814782  flash -0.03814782

逆伝播のために保存する量は、1ヘッドあたり次のように変わります（d=64。P と O は fp16、L は fp32 として計算）。

代わりに、逆伝播で QKᵀ をもう一度計算します。FLOPs はその分増えますが、n×n の行列を保存して HBM から読み直す必要がなくなります。前の節の数え上げのとおり、GPU ではこの読み書きの方が重くなりがちです。

精度：fp16 で計算したときの誤差の実測

この節では、GPU の FlashAttention が使う fp16・bf16 の入力で、どこを高い精度で計算すると誤差が小さくなるかを、numpy で再現した数字で分かります。

flash-attn の README によると、FlashAttention-2 が扱うデータ型は fp16 と bf16 です。行列の積は高い精度で足し合わせ、softmax の統計も高い精度で持つのが一般的な作りです。numpy で、fp16 の入力（n=1,024・d=64）に対して計算の精度を変え、同じ fp16 の入力を float64 で計算した値との差を測りました。

出力の値の大きさは最大 0.273、fp16 の機械イプシロンは 0.000977 です。fp32 で計算して最後に fp16 に丸めた場合の誤差（6.28e-05）は、0.273 付近の fp16 の刻み（2.44e-4）の半分以下で、出力を fp16 に丸める分に収まっています。全部 fp16 にすると誤差が約6倍になりました。和 ℓ と途中の出力 acc を fp16 で足し続けるためです。タイル化そのものは誤差を増やさず、誤差を決めるのは「どの値をどの精度で持つか」です。

FlashAttention-3 の論文は FP8 での誤差を扱っており、ブロックごとの量子化と、ランダムな符号を掛けたアダマール変換で外れ値をならす incoherent processing によって、ベースラインの FP8 の注意より誤差（RMSE）が2.6分の1（9.1e-3 と 2.4e-2）になったと報告しています（4.3節・表3）。

FlashAttention-2・FlashAttention-3 の改良点（論文の範囲）

この節では、FlashAttention-2 と FlashAttention-3 が何を変え、どれだけ速くなったと報告しているかが、原論文で確かめた範囲で分かります。

どちらも計算する値は同じ厳密な注意で、変えたのは GPU の上での仕事の分け方と実行のさせ方です。数字はすべて論文の報告（GPU）です。

1. 1FlashAttention：n×n の行列を HBM に書き出さないタイル化・オンライン softmax・逆伝播での再計算
2. 2FlashAttention-2：行列積以外の演算を減らし、並列に分ける最後に1回だけ割る・系列長方向の並列化・Q をワープに分ける
3. 3FlashAttention-3：データの移動と softmax を行列積に重ねるワープ特化・TMA・pingpong・2段のパイプライン・FP8

FlashAttention-2：行列積以外を減らす理由

FlashAttention-2 の論文の3.1節は、A100 の FP16/BF16 の行列積の性能が 312TFLOPs/s なのに対し、行列積以外の FP32 の演算は 19.5TFLOPs/s しかなく、「行列積以外の1 FLOP は行列積の1 FLOP の16倍高くつく」と述べています。この記事の Algorithm 1 のコードで毎回 ℓ を掛けて戻して割っていた部分が、まさにこの「行列積以外の演算」です。3.3節では、FlashAttention が K・V を4つのワープに分けていた方式（split-K）は、ワープ間で途中結果を共有メモリに書いて足し合わせる必要があったのに対し、FlashAttention-2 は Q を4つのワープに分けることで、順伝播でワープ間のやり取りを無くしたと説明しています。PyTorch の標準の注意に対しては最大10倍、Triton の FlashAttention の実装に対しては1.3〜2.5倍速いと報告しています（4.1節）。

FlashAttention-3：指数の計算が行列積に追いつかない問題

FlashAttention-3 の論文は、FlashAttention-2 が H100 で使えている性能は約35％にとどまると述べています。3.1節によると、H100 SXM5 の FP16 の行列積は 989TFLOPS なのに対し、指数などの特殊関数は 3.9TFLOPS しかなく、ヘッド次元128の FP16 の注意では指数の計算が行列積の50％ほどの時間をとりうるとしています。そこで、あるワープグループが softmax をしている間に別のワープグループが行列積をする pingpong スケジューリングと、1つのワープグループの中で行列積と softmax を2段に重ねるパイプラインを入れました。長い系列では FP16 で cuDNN を上回り、FP8 は同程度と報告しています。

実際に使うとき（flash-attn・PyTorch SDPA・Transformers）と判断表

この節では、FlashAttention を自分で書かずに使う3つの入口と、使う・使わないの判断の目安が、公式ドキュメントで確かめた範囲で分かります。手元に GPU が無いため、この節のコードは動かしていません。

flash-attn（Dao-AILab の公式実装）：README によると、FlashAttention-2 は CUDA 12.0 以降（または ROCm 6.0 以降）・PyTorch 2.2 以降が要り、GPU は Ampere・Ada・Hopper（A100、RTX 3090/4090、H100 など）、データ型は fp16 と bf16、ヘッド次元は256まで対応です。Turing（T4 など）は対象外です。入れ方は pip install flash-attn --no-build-isolation で、メモリが少なく CPU コアが多いマシンでは MAX_JOBS=4 のように並列のビルド数を絞るよう書かれています。FlashAttention-3 は Hopper 向けのベータ版で、リポジトリの hopper ディレクトリからビルドします。テンソルの並びは (batch, seqlen, nheads, headdim) です。

from flash_attn import flash_attn_func
# q, k, v: (batch, seqlen, nheads, headdim)、fp16 か bf16、CUDA 上
out = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=True)

PyTorch の scaled_dot_product_attention（SDPA）：PyTorch 2.15 の公式ドキュメントでは、実装として FlashAttention-2・Memory-Efficient Attention・C++ の標準的な実装の3つを挙げ、入力に応じて自動で選ぶとしています。テンソルの並びは (N, ..., ヘッド, L, E) で、flash-attn と順番が違います。実装を指定するには torch.nn.attention.sdpa_kernel() を使います。

import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend
# q, k, v: (batch, nheads, seqlen, headdim)
with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Hugging Face Transformers：公式ドキュメント（Attention Interface）では、from_pretrained() の attn_implementation に "flash_attention_2"・"flash_attention_3"・"sdpa" などを渡し、読み込んだ後も model.set_attn_implementation("sdpa") で切り替えられるとしています。推論サーバーで KV キャッシュと組み合わせる使い方は、関連記事「vLLMとは｜使い方・インストール・OpenAI互換APIとOllamaとの違い」で扱っています。

よくあるエラー・つまずき

この節では、FlashAttention を自作・利用するときに出やすいエラーと、手元で再現した原文・直し方が分かります。

2つの「結果がずれる」例は、n=1,024・d=64・B=128 の float64 で、正しい実装の誤差 3.05e-16 に対する値です。どちらもエラーも警告も出ないため、素朴な計算と比べるテストを最初に書いておくのが確実です。

よくある質問

この節では、FlashAttention の仕組みについてよく聞かれる質問への答えが分かります。

Q. FlashAttention は近似ですか。

近似ではありません。論文の題名どおり厳密な注意で、オンライン softmax は式の上で同じ値を出します。手元では float64 で 1e-15 以下まで一致し、float32 の誤差は素朴な計算と同じ水準でした。近似なのは、論文で拡張として示されたブロックスパース版です。

Q. 計算量は減るのですか。

FLOPs は O(N²d) のままで、逆伝播では再計算の分だけ増えます。減るのは HBM への読み書きと、n×n の行列のためのメモリです。

Q. メモリが「線形」になるとはどういう意味ですか。

入力と出力のほかに必要なメモリが系列長 n に比例する、という意味です（論文の定理1）。手元では出力を除いた追加分が n=512 で 0.098MiB、n=16,384 で 0.159MiB で、増えた分は logsumexp の n 個分でした。

Q. オンライン softmax は FlashAttention の発明ですか。

softmax の正規化の項を1パスで求める方法は Milakov・Gimelshein（2018）が示しています。Rabe・Staats（2021）は、注意を O(n²) のメモリなしで計算できることを示し、系列長16,384 で自己注意のメモリのオーバーヘッドを推論で59分の1にしたと報告しています。FlashAttention は、これを GPU のメモリ階層に合わせて HBM の読み書きを減らす形に組み、IO の量を解析した点が新しいところです。

Q. CPU でも FlashAttention で速くなりますか。

手元の numpy では、ブロックを128〜256にしてようやく素朴な計算と同じくらいで、それより速くはなりませんでした。numpy では1つのブロックの計算を高速なメモリの中で融合できないためです。メモリを減らす効果は CPU でもそのまま得られます。

Q. ブロックの大きさはどう決めますか。

論文の Algorithm 1 は SRAM の大きさ M から Bc=⌈M/4d⌉、Br=min(⌈M/4d⌉, d) と決めています。Q・K・V・O のブロックが SRAM に収まる大きさです。実際のカーネルはヘッド次元ごとに調整された値を使います。

Q. 推論（1トークンずつの生成）でも効きますか。

1トークンずつの生成では Q が1行だけになり、n×n の行列はもともと作られません。この場合は過去の K・V をどう持つかが問題になり、本サイトで公開予定の「LLM推論の高速化｜KVキャッシュと投機的デコーディングを自作で実測」で扱っています。長い K・V を分けて計算し、部分結果を合わせる場合は、この記事の2ブロックの結合の式がそのまま使えます。

Q. FlashAttention-2 と 3 のどちらを使うべきですか。

FlashAttention-3 は Hopper（H100 など）向けに作られており、README ではベータ版の扱いです。Ampere・Ada の GPU では FlashAttention-2 が対象です。どちらも PyTorch の SDPA や Transformers の attn_implementation から選べるため、手元の GPU で両方を測って決めるのが確実です。

Q. Attention の計算そのものから復習したいときは。

Q・K・V と softmax の意味は、関連記事「ニューラルネットワークとは｜仕組み・学習・AttentionをNumPyで実演」で扱っています。この記事のコードの配列操作（keepdims・ブロードキャスト）は関連記事「NumPyの使い方 基本一覧｜配列の作成・演算・形状変更」で扱っています。

Q. GPU のメモリが足りるかを見積もるには。

モデルの重み・KV キャッシュ・活性化の見積もり方は関連記事「VRAMとは｜確認方法・必要量の目安とLLMの計算式をPythonで解説」で扱っています。LLM の全体の仕組みは関連記事「生成AI（LLM）の仕組み8選｜ハルシネーション・トークン」で扱っています。

まとめ

この節では、記事全体の要点が分かります。

1. FlashAttention は厳密な注意を、n×n の行列を書き出さずにブロックごとに計算する。手元の numpy で、float64 で 1e-15 以下、float32 で素朴な計算と同じ水準で一致した。
2. 鍵はオンライン softmax。最大値 m と和 l を持ち回り、最大値が変わったら exp(m旧−m新) で ℓ と途中の出力の両方を補正する。結合は順番に依らない。
3. メモリは n² から n へ。n=4,096 で素朴な 192MiB に対し、出力を除く追加分は 0.11MiB。逆伝播も P を保存せず logsumexp から作り直す。
4. 速くなるのは計算ではなく HBM の読み書きが減るから。論文は Θ(Nd+N²) を Θ(N²d²/M) にすると証明。numpy では融合ができず、手元では素朴な計算と同じくらいの速さだった。
5. FlashAttention-2 は行列積以外の演算と並列化、FlashAttention-3 は Hopper の非同期実行と FP8 を改良。因果マスクではブロックの約半分を飛ばせ、行が全部マスクのときの NaN に注意する。

参考・出典

この節では、この記事で根拠にした論文と公式ドキュメントが分かります（確認日 2026年11月4日）。

• Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré「FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness」2022：https://arxiv.org/abs/2205.14135
• Tri Dao「FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning」2023：https://arxiv.org/abs/2307.08691
• Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao「FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision」2024：https://arxiv.org/abs/2407.08608
• Maxim Milakov, Natalia Gimelshein「Online normalizer calculation for softmax」2018：https://arxiv.org/abs/1805.02867
• Markus N. Rabe, Charles Staats「Self-attention Does Not Need O(n²) Memory」2021：https://arxiv.org/abs/2112.05682
• Dao-AILab/flash-attention（GitHub の README）：https://github.com/Dao-AILab/flash-attention
• PyTorch 2.15 公式ドキュメント「torch.nn.functional.scaled_dot_product_attention」：https://docs.pytorch.org/docs/2.15/generated/torch.nn.functional.scaled_dot_product_attention.html
• Hugging Face Transformers 公式ドキュメント「Attention Interface」：https://huggingface.co/docs/transformers/main/en/attention_interface

株式会社RJC ― SI事業・SES事業・AI駆動開発。RJCは一緒に成長を楽しめる会社です。