この記事の対象読者

PyTorch で softmax(Q @ K.T / sqrt(d)) @ V を自分で書いたことがある人

前提知識はこの1行だけです。Attentionの仕組みそのものも本文で1から定義し直します。

この記事で得られること

読み終えたあと、手元の手書きAttentionを F.scaled_dot_product_attention に置き換えてよいか、置き換えるならマスク・dropout・因果マスクの3つの引数をどう書き換えるかを決められます。

置き換えの判断材料として、手書き版との数値差、メモリ、速度の実測値
置き換えで静かに壊れる3箇所と、それぞれを塞ぐ書き方
そのまま貼れる置き換え済みのAttentionモジュール

この記事で扱わないこと

FlashAttentionのカーネル内部の詳細。考え方だけ触れ、詳細は関連記事に回します
GPU上での速度比較。今回の計測はすべてCPUです
FlexAttentionなど、SDPAの外側にある新しいAPI

なぜこれを書くのか

小型GPTをスクラッチで組む教材を作ったとき、Attention は教科書どおり手で書きました。教材としてはそれで正解でした。ところが Hugging Face の Transformers を読むと attn_implementation="sdpa" が当たり前のように並んでいて、自分の手書き版をこの1関数に置き換えてよいのか、置き換えたら何が変わるのかを、私は説明できませんでした。それが悔しくて、公式ドキュメントを頭から読み、手元で測りました。

読者に考えてほしいのは速さの話より、同じ式を計算する関数でも引数の意味が違えば静かに壊れるという点です。置き換えの前に等価性を測る習慣を、自分のコードベースでどう回すか。その材料として読んでください。

計測環境はGPUのないクラウドの作業環境です。Xeon 2コア、メモリ約8GBと貧弱ですが、そのぶん手書き版が8GiBの確保に失敗して落ちる境界がはっきり見えました。

項目	値
PyTorch	2.14.1
Python	3.13.17
CPU	Intel Xeon 2コア
メモリ	約8GB
GPU	なし

1. SDPAとは何か ── 手書きAttentionと同じ式を計算する関数

このセクションで分かること：SDPAが何を計算する関数で、手書き版と結果がどこまで一致するか

図書館のレファレンス係で考える

この記事では、Attentionを図書館のレファレンス係にたとえます。利用者が質問カードを持ってくると、係は蔵書の背表紙と1冊ずつ照らし合わせ、関係が深そうな本の中身を割合を決めて混ぜ合わせ、答えとして返します。この作業1回分がAttentionです。

記号は使う直前に1つずつ導入します。

$Q$ はクエリで、質問カードの束です。型は4次元テンソルで、形は $(N, H_q, L, E)$ です
$L$ は質問カードの枚数で、整数です。公式ドキュメントでは target sequence length と呼ばれます
$E$ はカード1枚に書ける項目数で、整数です。埋め込み次元とも呼びます
$K$ はキーで、本の背表紙の束です。形は $(N, H, S, E)$ です
$S$ は蔵書の冊数で、整数です。source sequence length と呼ばれます
$V$ はバリューで、本の中身の束です。形は $(N, H, S, E_v)$ です
$B$ はバイアスで、本に貼る閲覧札です。形は $(L, S)$ の行列で、値は閲覧可なら $0$、閲覧禁止なら $-\infty$ です

道具がそろったので式を書きます。

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{E}} + B\right) V

この式が言っているのは、次の4段階です。

$QK^{\top}$ で、全質問カード×全背表紙の相性表を作る。大きさは $L \times S$
$\sqrt{E}$ で割って、項目数が多いほど相性点が膨らむ偏りを打ち消す
$B$ を足して、閲覧禁止の本の相性点を $-\infty$ にする
softmaxで相性点を「どの本を何割読むか」の配分に変え、その割合で $V$ を混ぜる

記号	読み	図書館での役割	型
$Q$	クエリ	質問カードの束	テンソル
$K$	キー	本の背表紙	テンソル
$V$	バリュー	本の中身	テンソル
$L$	ターゲット長	質問カードの枚数	整数
$S$	ソース長	蔵書の冊数	整数
$E$	埋め込み次元	カード1枚の項目数	整数
$B$	バイアス	閲覧札	行列
$H$	ヘッド数	並行して働く係員の人数	整数

この式をそのまま1関数にしたのが torch.nn.functional.scaled_dot_product_attention、略してSDPAです。シグネチャは次のとおりで、scale だけはキーワード引数でしか渡せません。

F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
                               is_causal=False, scale=None, enable_gqa=False)

本当に同じ答えが返るのか測る

公式ドキュメントは、この関数が上の式と等価な効率的実装だと説明しています。言葉だけでは置き換えの判断材料にならないので、手書き版と突き合わせました。

check_equiv.py

import math
import torch
import torch.nn.functional as F


def naive_attention(q, k, v, causal=False):
    """教科書どおりの手書きAttention"""
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    if causal:
        L, S = q.size(-2), k.size(-2)
        mask = torch.ones(L, S, dtype=torch.bool).tril()
        scores = scores.masked_fill(~mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v


torch.manual_seed(0)
for dtype in (torch.float32, torch.float64):
    for causal in (False, True):
        q, k, v = (torch.randn(2, 8, 512, 64, dtype=dtype) for _ in range(3))
        ref = naive_attention(q, k, v, causal=causal)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
        diff = (ref - out).abs().max().item()
        print(f"{str(dtype):15s} causal={causal!s:5s} max|diff|={diff:.2e}")

dtype	因果マスク	手書き版との最大絶対誤差
float32	なし	4.47e-07
float32	あり	5.36e-07
float64	なし	7.22e-16
float64	あり	1.11e-15

float32で $10^{-7}$、float64で $10^{-15}$ という差は、どちらもその型の丸め誤差の水準です。SDPAは手書き版と同じ答えを返すと言ってよい結果でした。

2. 同じ式なのに何が違うのか ── 相性表を机に広げない

このセクションで分かること：SDPAに置き換えるとメモリと時間がどれだけ変わるか

前節で、SDPAの出力は手書き版と丸め誤差の範囲で一致すると確かめました。答えが同じなら、違いは計算の途中にしかありません。そこで、途中で使うメモリと時間を測ります。

手書き版は、相性表 $QK^{\top}$ を丸ごとメモリに置きます。図書館で言えば、全質問カード×全蔵書の表を机いっぱいに広げてから作業を始めるやり方です。表の大きさはこう書けます。新しい記号はありません。

\text{相性表のバイト数} = H \times L \times S \times 4

この式が言っているのは、係員 $H$ 人がそれぞれ $L \times S$ マスの表を持ち、1マスがfloat32の4バイトを食う、ということです。$L$ と $S$ を2倍にすると表は4倍になります。

一方、SDPAが選ぶ融合カーネル（複数の演算を1つにまとめ、中間結果をメモリに書き出さない実装）は、表を広げません。棚の区画ごとに背表紙を見て回り、その時点の最大相性点と配分の合計だけを手元のメモに書き足していきます。最後にメモから辻褄を合わせるので、机は小さいまま済みます。この考え方の詳細は Flash Attention の記事に書きました。

bench_mem.py

import subprocess
import sys

CHILD = """
import math, resource, sys, time, torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
impl, L = sys.argv[1], int(sys.argv[2])
q, k, v = (torch.randn(1, 8, L, 64) for _ in range(3))
base = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
t0 = time.perf_counter()
with torch.no_grad():
    if impl == "naive":
        out = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(64), -1) @ v
    elif impl == "math":
        with sdpa_kernel(SDPBackend.MATH):
            out = F.scaled_dot_product_attention(q, k, v)
    else:
        out = F.scaled_dot_product_attention(q, k, v)
dt = time.perf_counter() - t0
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f"{(peak - base) / 1024:.0f}MB/{dt:.2f}s")
"""

IMPLS = ("naive", "math", "sdpa")
print("| L | " + " | ".join(IMPLS) + " |")
for L in (1024, 2048, 4096, 8192, 16384):
    cells = []
    for impl in IMPLS:
        r = subprocess.run([sys.executable, "-c", CHILD, impl, str(L)],
                           capture_output=True, text=True)
        cells.append(r.stdout.strip() if r.returncode == 0 else "確保失敗")
    print(f"| {L} | " + " | ".join(cells) + " |")

1回ずつ別プロセスで実行し、最大常駐メモリの増分を測っています。$H=8$、$E=64$、float32、推論モードです。

L	手書き	SDPA・MATH強制	SDPA・指定なし
1024	116MB / 0.03s	99MB / 0.04s	19MB / 0.01s
2048	273MB / 0.13s	316MB / 0.16s	20MB / 0.03s
4096	1043MB / 0.49s	1190MB / 0.63s	25MB / 0.12s
8192	4117MB / 2.14s	4664MB / 2.62s	33MB / 0.49s
16384	確保失敗	確保失敗	49MB / 2.04s

読み取れることは3つです。

手書き版の $L=8192$ は4117MBでした。式で出す相性表は $8 \times 8192^2 \times 4$ バイトで2048MiBなので、相性表とsoftmax後の表の2枚分がちょうど載っています
$L=16384$ の手書き版は、you tried to allocate 8589934592 bytes で落ちました。$8 \times 16384^2 \times 4 = 8{,}589{,}934{,}592$ で、上の式とバイト単位まで一致します
指定なしのSDPAは、$L$ を16倍にしても19MBから49MBにしか増えず、速度も手書き版の約4倍でした

予想外だったのは真ん中の列です。同じSDPAでも、MATHバックエンドを強制すると手書き版より重くなりました。SDPAという関数名ではなく、中で誰が計算しているかが性能を決めています。

GPU でも事情は同じで、手書き版の相性表は VRAM を $L^2$ で食います。

3. 中で誰が計算しているのか ── 4つのバックエンドと自動選択

このセクションで分かること：SDPAの中身を選ぶ仕組みと、自分の環境で何が選ばれたかを確かめる方法

前節で、同じSDPAでもMATHを強制すると手書き版より重くなることを見ました。つまり置き換えの効果は、SDPAがどの実装を選ぶかで決まります。ここではその選び方を確かめます。

SDPAは、呼ばれるたびに入力の形・型・デバイスを見て、条件を満たす実装の中から最適なものを自動で選びます。この内部の実装をバックエンドと呼びます。図書館で言えば、同じ窓口の奥に作業スタイルの違う係員チームが控えていて、質問の種類に応じて担当が割り当てられる仕組みです。

バックエンド	中身	図書館での作業スタイル
`FLASH_ATTENTION`	FlashAttention-2系の融合カーネル	区画ごとに巡回してメモだけ持つ
`EFFICIENT_ATTENTION`	xFormers由来のメモリ効率カーネル	同じく表を広げない別チーム
`CUDNN_ATTENTION`	NVIDIA cuDNN の実装	外部から来た専門チーム
`MATH`	式をそのままPyTorchの演算で組んだもの	表を机いっぱいに広げる

候補を絞りたいときは torch.nn.attention.sdpa_kernel コンテキストマネージャを使います。CPUで1つずつ強制してみました。

check_backends.py

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

BACKENDS = [
    SDPBackend.MATH,
    SDPBackend.FLASH_ATTENTION,
    SDPBackend.EFFICIENT_ATTENTION,
    SDPBackend.CUDNN_ATTENTION,
]

q, k, v = (torch.randn(1, 8, 256, 64) for _ in range(3))

for backend in BACKENDS:
    try:
        with sdpa_kernel(backend):
            F.scaled_dot_product_attention(q, k, v)
        result = "OK"
    except RuntimeError as e:
        result = "NG: " + str(e).splitlines()[0][:60]
    print(f"{backend.name:20s} {result}")

バックエンド	CPUでの結果
MATH	動いた
FLASH_ATTENTION	動いた
EFFICIENT_ATTENTION	No viable backend for scaled_dot_product_attention was found
CUDNN_ATTENTION	同上

ここで2つ目の驚きがありました。公式ドキュメントのNoteには、最適化カーネルを呼ぶのはCUDAバックエンドで、それ以外ではPyTorch実装を使う、という趣旨の記述があります。ところがCPUでも FLASH_ATTENTION が通り、前節の49MBという数字もMATHでは説明がつきません。そこでプロファイラで、実際に走った演算名を見ました。

check_dispatch.py

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.profiler import ProfilerActivity, profile

q, k, v = (torch.randn(1, 8, 1024, 64) for _ in range(3))


def called_ops(backend=None):
    """SDPA 1回の呼び出しで実際に走った aten 演算名を返す"""
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        if backend is None:
            F.scaled_dot_product_attention(q, k, v)
        else:
            with sdpa_kernel(backend):
                F.scaled_dot_product_attention(q, k, v)
    keys = {e.key for e in prof.key_averages()}
    return sorted(k for k in keys if "attention" in k or "softmax" in k)


print("指定なし:", called_ops())
print("MATH    :", called_ops(SDPBackend.MATH))

呼び方	実際に走った主な演算
指定なし	`aten::_scaled_dot_product_flash_attention_for_cpu`
MATH強制	`aten::_scaled_dot_product_attention_math` と `aten::_safe_softmax`

PyTorch 2.14.1のCPUでは、指定なしでCPU向けのFlash系カーネルが選ばれていました。ドキュメントの記述より実装が先に進んでいるようです。自分の環境で何が選ばれるかは、ドキュメントを読むよりプロファイラで見るほうが確実です。MATHのほうに出てくる _safe_softmax は、第4節で効いてきます。

選択の流れをまとめるとこうなります。

GPUで特定の融合カーネルを使いたい場合、sdpa_kernel でMATHを候補から外しておくと、条件を満たさないときに黙ってMATHへ落ちるのを防げます。落ちた理由は警告として表示されます。

ここまでのまとめ

節	分かったこと
1	SDPAは手書き版と丸め誤差の範囲で同じ答えを返す
2	融合カーネルが選ばれれば、メモリは $L^2$ で増えず、CPUでも約4倍速い
3	何が選ばれるかは入力次第で、MATHに落ちると手書き版より重い

ここまでで、置き換える価値はあると言えます。ただし第1節で一致を確かめたのは、マスクなしと単純な因果マスクの2通りだけでした。実際のモデルには、パディングのマスク、dropout、推論時の KVキャッシュが絡みます。次節はそこを測ります。

4. 置き換えで静かに壊れる3つの罠

このセクションで分かること：手書き版の感覚のまま引数を渡すと結果が変わる3箇所と、その直し方

前節までで、式は同じ・速い・自動で最適な実装が選ばれる、と確かめました。しかし式が同じでも、引数の意味が手書き版や nn.MultiheadAttention の流儀と違う箇所が3つあります。どれもエラーは出ず、数字だけが静かにずれます。

罠1. ブールマスクのTrueの意味が逆

SDPAの attn_mask にブールマスクを渡すと、Trueは「このペアをAttentionに使う」を意味します。一方、nn.MultiheadAttention の key_padding_mask では、Trueは「パディングなので捨てる」です。公式ドキュメントのNoteも、MHAから移行するならブールマスクを反転せよと注意しています。図書館で言えば、同じ赤い札が片方の館では閲覧可、もう片方の館では閲覧禁止を意味している状態です。

pitfall_mask.py

import torch
import torch.nn.functional as F

torch.manual_seed(0)
B, H, L, D = 1, 1, 4, 8
q, k, v = (torch.randn(B, H, L, D) for _ in range(3))

# 最後の2トークンが padding だとする
is_pad = torch.tensor([[False, False, True, True]])        # MHA の流儀 True=捨てる
keep = ~is_pad                                              # SDPA の流儀 True=使う

# 正解: padding を使わずに計算した結果
ref = F.scaled_dot_product_attention(q, k[..., :2, :], v[..., :2, :])

# 罠: MHA の key_padding_mask をそのまま渡す
bad = F.scaled_dot_product_attention(q, k, v, attn_mask=is_pad[:, None, None, :])
# 正: 反転してから渡す
ok = F.scaled_dot_product_attention(q, k, v, attn_mask=keep[:, None, None, :])

print("そのまま渡す  max|diff| =", f"{(bad - ref).abs().max().item():.2e}")
print("反転して渡す  max|diff| =", f"{(ok - ref).abs().max().item():.2e}")

# おまけ: 全部 False の行があると何が出るか
all_false = torch.zeros(B, H, L, L, dtype=torch.bool)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=all_false)
print("全行マスク時の出力に NaN を含む:", torch.isnan(out).any().item())
print("全行マスク時の出力の最大絶対値:", out.abs().max().item())

渡し方	正解との最大絶対誤差
MHAの流儀のまま渡す	3.29e+00
反転して渡す	5.96e-08

そのまま渡すと、パディングだけを読んで本物のトークンを無視する計算になります。反転すれば一致します。

おまけの結果も書いておきます。全部の本が閲覧禁止の行、つまり $B$ の1行がすべて $-\infty$ の行では、手書き版のsoftmaxはNaNを返します。SDPAはCPUの指定なしでもMATH強制でも、NaNではなく0を返しました。MATH経路では、第3節で見た _safe_softmax がこの安全側の挙動を担っていると考えられます。左詰めパディングで先頭行が全部マスクされても、SDPAならNaNが伝播しにくいということです。

罠2. eval()にしてもdropoutが止まらない

nn.Dropout は model.eval() で自動的に止まります。SDPAの dropout_p は止まりません。公式ドキュメントの警告はこう書いています。

This function always applies dropout according to the specified dropout_p argument.
和訳：この関数は、指定された dropout_p 引数に従って常にdropoutを適用します。

図書館で言えば、研修用の「ランダムに本を隠して係員を鍛える訓練」を、本番の窓口でも続けてしまう状態です。

pitfall_dropout.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class AttnBad(nn.Module):
    def __init__(self, p=0.1):
        super().__init__()
        self.p = p

    def forward(self, q, k, v):                     # 罠: 常に self.p を渡す
        return F.scaled_dot_product_attention(q, k, v, dropout_p=self.p)


class AttnGood(AttnBad):
    def forward(self, q, k, v):                     # 正: 学習中だけ p を渡す
        p = self.p if self.training else 0.0
        return F.scaled_dot_product_attention(q, k, v, dropout_p=p)


q, k, v = (torch.randn(1, 8, 128, 64) for _ in range(3))
for cls in (AttnBad, AttnGood):
    m = cls().eval()                                # 推論モードにしたつもり
    with torch.no_grad():
        a, b = m(q, k, v), m(q, k, v)
    print(f"{cls.__name__:9s} 同じ入力で2回推論した差 max|diff| = "
          f"{(a - b).abs().max().item():.2e}")

実装	eval後、同じ入力で2回推論した差
`dropout_p=self.p` 固定	7.07e-01
`self.training` で分岐	0.00e+00

torch.no_grad() を付けても止まりません。推論のたびに出力が変わるのに、エラーも警告も出ないのがこの罠の怖いところです。

罠3. KVキャッシュ推論でis_causal=Trueを付けると最初の1冊しか読めない

文章生成の推論では、過去トークンの $K$ と $V$ をキャッシュに溜め、新しい1トークン分の $Q$ だけを計算します。このとき $L=1$、$S$ は過去の長さで、$L \ne S$ になります。

is_causal=True の因果マスク（未来のトークンを見せないための下三角マスク）は、$L \ne S$ のとき左上揃えで作られます。$L=1, S=6$ なら、閲覧可は最初の1冊だけです。本当に欲しいのは右下揃え、つまり最新の質問カードが全蔵書を読める形です。

pitfall_causal.py

import torch
import torch.nn.functional as F
from torch.nn.attention.bias import causal_lower_right

torch.manual_seed(0)
S = 6                                   # KVキャッシュに溜まった長さ
q_all, k, v = (torch.randn(1, 1, S, 8) for _ in range(3))

# 正解: S トークン全部を一度に因果マスク付きで計算し、最後の行だけ取る
ref = F.scaled_dot_product_attention(q_all, k, v, is_causal=True)[..., -1:, :]

q_last = q_all[..., -1:, :]             # 推論時は最新の1トークンだけが query

# 罠: L=1, S=6 で is_causal=True
bad = F.scaled_dot_product_attention(q_last, k, v, is_causal=True)
# 正: 1トークン decode なら因果マスク不要
ok1 = F.scaled_dot_product_attention(q_last, k, v, is_causal=False)
# 正: 右下揃えの因果バイアスを渡す
ok2 = F.scaled_dot_product_attention(q_last, k, v, attn_mask=causal_lower_right(1, S))

for name, out in [("is_causal=True", bad), ("is_causal=False", ok1),
                  ("causal_lower_right", ok2)]:
    print(f"{name:20s} max|diff| vs 正解 = {(out - ref).abs().max().item():.2e}")

# 何が起きているか: upper-left 揃えの L=1 マスクは key 0 しか見せない
print(torch.ones(1, S, dtype=torch.bool).tril())

渡し方	正解との最大絶対誤差
`is_causal=True`	3.15e+00
`is_causal=False`	1.19e-07
`attn_mask=causal_lower_right(1, S)`	1.19e-07

最後の print は [[True, False, False, False, False, False]] を出します。最新トークンが、文頭の1トークンしか見ていません。1トークンずつ生成するなら is_causal=False で足ります。複数トークンをまとめて処理するチャンク推論では、右下揃えの causal_lower_right(L, S) を使います。

この罠は学習時には絶対に出ません。学習では $L = S$ なので、左上揃えと右下揃えが一致するからです。推論だけ品質が落ちる、という形で後から見つかります。

5. 置き換えの実装 ── 3つの罠を塞いだモジュール

このセクションで分かること：手書きAttentionからの具体的な書き換え方と、そのまま使えるモジュール

前節で見つけた3つの罠は、どれも引数1つで塞げます。そこで、第1節の手書き版をSDPAに置き換えるときの差分から見ます。

- scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
- scores = scores.masked_fill(is_pad[:, None, None, :], float("-inf"))
- scores = scores.masked_fill(~causal_tri, float("-inf"))
- w = torch.softmax(scores, dim=-1)
- w = self.attn_drop(w)                  # nn.Dropout は eval で止まる
- y = w @ v
+ keep = ~is_pad[:, None, None, :]       # 罠1: True=使う に反転
+ y = F.scaled_dot_product_attention(
+     q, k, v, attn_mask=keep & causal_tri,
+     dropout_p=self.p if self.training else 0.0)  # 罠2

パディングと因果マスクを両方使う場合は、2つをANDで合成して attn_mask に渡します。attn_mask と is_causal=True を同時に渡すとエラーになるためです。

これを、ヘッド分割と出力射影まで含めたモジュールにしました。Query ヘッド数と Key/Value ヘッド数を別にすると、GQA（Grouped Query Attention：複数のQueryヘッドが1組のKey/Valueヘッドを共有する方式）にもなります。図書館で言えば、32人の係員が8つの棚を4人ずつ共有する形です。SDPAでは enable_gqa=True を渡すと、棚の共有をSDPA側がやってくれます。

sdpa_attention.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model=512, n_head=8, n_kv_head=8, p_drop=0.1):
        super().__init__()
        assert d_model % n_head == 0 and n_head % n_kv_head == 0
        self.h, self.kvh, self.p = n_head, n_kv_head, p_drop
        self.dh = d_model // n_head
        self.q_proj = nn.Linear(d_model, n_head * self.dh)
        self.kv_proj = nn.Linear(d_model, 2 * n_kv_head * self.dh)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, pad_mask=None):
        """x: (N, L, d_model) / pad_mask: (N, L) で True=パディング"""
        N, L, _ = x.shape
        q = self.q_proj(x).view(N, L, self.h, self.dh).transpose(1, 2)
        kv = self.kv_proj(x).view(N, L, 2, self.kvh, self.dh)
        k, v = kv.unbind(2)
        k, v = k.transpose(1, 2), v.transpose(1, 2)

        mask, causal = None, True
        if pad_mask is not None:                       # 罠1: True=使う へ反転
            keep = ~pad_mask[:, None, None, :]         # (N, 1, 1, S)
            tri = torch.ones(L, L, dtype=torch.bool, device=x.device).tril()
            mask, causal = keep & tri, False           # is_causal と併用しない
        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=mask, is_causal=causal,
            dropout_p=self.p if self.training else 0.0,  # 罠2
            enable_gqa=self.h != self.kvh)
        return self.out(y.transpose(1, 2).reshape(N, L, -1))

このモジュールは学習と全文一括の推論を想定しています。KVキャッシュで1トークンずつ生成する経路を足すときは、罠3のとおり is_causal=False か causal_lower_right に切り替えてください。

動作確認として、左詰めパディングを付けた入力と、パディングを最初から持たない入力の出力を比べました。

test_module.py

import torch
from sdpa_attention import CausalSelfAttention

torch.manual_seed(0)
for kvh in (8, 2):
    m = CausalSelfAttention(d_model=64, n_head=8, n_kv_head=kvh).eval()
    x = torch.randn(2, 10, 64)
    pad = torch.zeros(2, 10, dtype=torch.bool)
    pad[:, :3] = True                     # 左詰め: 先頭3トークンが padding
    with torch.no_grad():
        full = m(x, pad_mask=pad)[:, 3:]
        trunc = m(x[:, 3:])               # padding を最初から持たない入力
        nomask = m(x)[:, 3:]              # マスクを付け忘れた場合
        a, b = m(x), m(x)
    print(f"kv_head={kvh}: マスクあり {(full - trunc).abs().max():.1e}, "
          f"マスクなし {(nomask - trunc).abs().max():.1e}, "
          f"eval再現性 {(a - b).abs().max():.1e}")

KVヘッド数	マスクあり	マスク付け忘れ	eval後の再現性
8	0.0e+00	8.4e-01	0.0e+00
2 (GQA)	0.0e+00	7.1e-01	0.0e+00

マスクを正しく渡せばパディングの有無で出力が変わらず、eval後は毎回同じ出力が返ります。GQAでも同じです。

Hugging Face Transformersで使う場合

既存のモデルを使う側なら、自分でモジュールを書く必要はありません。読み込み時に指定するだけです。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("モデル名", attn_implementation="sdpa")

Hugging Faceのドキュメントでは、sdpa はPyTorch組み込みのSDPA実装を使う選択肢として説明されています。第3節の sdpa_kernel で囲めば、生成時に使うカーネルも絞れます。

置き換えてよいかの判定

6. トラブルシューティング

実際に踏んだエラーと、出力が静かにずれる症状をまとめます。

症状・エラー	原因	対処
`Explicit attn_mask should not be set when is_causal=True`	マスクと `is_causal=True` の併用	因果マスクを自分で作り、ANDで合成して `attn_mask` だけ渡す
`takes from 3 to 6 positional arguments but 7 were given`	`scale` を位置引数で渡した	`scale=0.125` のようにキーワードで渡す
`Expected attn_mask dtype to be bool or float or to match query dtype`	int64やfloat64のマスク	`bool` にするか、`query` と同じ浮動小数型にする
`The size of tensor a (32) must match the size of tensor b (8)`	QとKVのヘッド数が違うのにGQAを有効にしていない	`enable_gqa=True` を渡す
`No viable backend for scaled_dot_product_attention was found`	`sdpa_kernel` で、その環境で動かないバックエンドだけを指定した	候補にMATHを含めるか、条件を満たす入力にする
パディング付きで出力がおかしい	ブールマスクのTrueの意味が逆	`~mask` で反転
eval後も推論のたびに出力が変わる	`dropout_p` を固定値で渡している	`self.p if self.training else 0.0`
学習は良いのに生成だけ品質が低い	$L \ne S$ で `is_causal=True`	`is_causal=False` か `causal_lower_right`
長い入力でメモリが足りない	MATHが選ばれて相性表を広げている	プロファイラで演算名を確認し、融合カーネルの条件を満たす

用語集

用語	意味	図書館での対応
SDPA	Scaled Dot Product Attention。PyTorchでは `F.scaled_dot_product_attention`	レファレンス係の窓口そのもの
融合カーネル	複数の演算を1つの実装にまとめ、中間結果をメモリに書き出さない計算	表を広げず巡回しながらメモだけ取る作業
バックエンド	SDPAの内部で実際に計算する実装	窓口の奥に控える係員チーム
ブールマスク	True/Falseで、どのペアを使うかを指定するマスク	閲覧札
因果マスク	未来のトークンを見せないための下三角マスク	まだ入荷していない本を閲覧禁止にする札
左上揃え・右下揃え	$L \ne S$ のとき下三角をどの角に合わせるか	最初の本から数えるか、最新の本から数えるか
GQA	Grouped Query Attention。複数のQヘッドが1組のKVヘッドを共有する方式	複数の係員が1つの棚を共有する
KVキャッシュ	推論で過去トークンの $K$, $V$ を保存して使い回す仕組み	過去の背表紙と中身の写しを手元に残しておく

学習ロードマップ

関連記事

Self-Attention：本記事の式の土台です
Transformer：SDPAが組み込まれる全体構造
Attention と KVキャッシュ：罠3の背景
Flash Attention：第2節で触れた「表を広げない」仕組みの詳細
CUDA：GPUでSDPAの融合カーネルを動かすための前提

参考文献

公式ドキュメントはPyTorch 2.14版を参照しています。

PyTorch公式APIリファレンス（SDPA本体）。引数の定義、dropoutの警告、MHAとのマスクの違い、バックエンドの説明が載っています。

PyTorch公式APIリファレンス（sdpa_kernel）。バックエンドを絞るコンテキストマネージャの使い方です。

PyTorch公式APIリファレンス（CausalVariant）。左上揃えと右下揃えの因果マスクを行列の例つきで定義しています。

(Beta) Implementing High-Performance Transformers with Scaled Dot Product Attention (SDPA)
邦題訳：SDPAによる高性能Transformerの実装（ベータ版）。PyTorch公式チュートリアルです。

Tri Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning", 2023
邦題訳：FlashAttention-2 ── より良い並列化と作業分割による高速なAttention。SDPAの FLASH_ATTENTION バックエンドの元になった論文です。

Hugging Face Transformers公式ドキュメント（Attention backends）。attn_implementation="sdpa" の指定方法が載っています。

まとめ

SDPAは、手書きAttentionと同じ式を計算する1関数です。今回の計測では、手書き版との差は丸め誤差の範囲で、融合カーネルが選ばれればメモリは $L^2$ で増えず、CPUでも約4倍速くなりました。手書き版が8GiBの確保に失敗した長さでも、SDPAは49MBで計算を終えています。

一方で、置き換えには3つの罠がありました。ブールマスクのTrueの意味が逆、dropout_p がevalで止まらない、$L \ne S$ の is_causal が左上揃え。どれもエラーにならず数字だけがずれます。図書館のたとえで言えば、窓口を新しくするのは正しい判断ですが、閲覧札の色の意味と、研修の続け方と、新着の数え方だけは引き継ぎ書を読み直す必要がある、ということです。

個人的に一番効いたのは、ドキュメントに書かれていない挙動、つまりCPUでもFlash系カーネルが選ばれることや、全マスク行がNaNでなく0になることを、プロファイラと数行のテストで確かめられた点でした。置き換える前に、第1節のような等価性チェックを自分の入力で1回回してください。それが一番安い保険です。

著者X

関連記事