この記事で決められること
読み終えたあと、推論で nn.MultiheadAttention をそのまま使うか、F.scaled_dot_product_attention で組んだ自前実装に切り替えるかを、自分の環境で測って決められるようになります。
対象読者
- PyTorch の nn.MultiheadAttention を使ったことがあり、batch_first やマスクの基本を知っている人
この記事で得られること
- 呼び方によって、どの内部経路が選ばれるかを確かめるスクリプト
- CPU で fastpath より SDPA 経路のほうが速かったという実測と、その条件
- nn.MultiheadAttention と数値が一致する SDPA 製の自前実装と、GQA への拡張
- 使い続けるか自前で組むかの判定フロー
この記事で扱わないこと
- GPU での計測と、NestedTensor でパディングを省いたときの効果
- MHA の数式の詳しい導出(前編で扱っています)
なぜこれを書くのか
前編で nn.MultiheadAttention の中身を手で組み直したとき、ソースの fastpath 判定に、ドキュメントの条件一覧にはない「num_heads が偶数」という条件を見つけました。fastpath は推論を速くするための経路だと説明されています。それなら実際どれだけ速いのかと思って測ったところ、手元では逆に遅いという結果が出ました。
「高速経路」と名前が付いているものを、測らずに信じてよいのか。この記事ではその問いを、数字と一緒に考えてもらえたらうれしいです。検証は GPU なしの Xeon 2コア・メモリ8GB の Linux 環境で、Python 3.13 と PyTorch 2.14.0 を CPU 実行しています。
1. MHAのおさらいと、2つの高速経路
最初に、この記事だけで話が通じるよう、MHA とたとえを短く定義しておきます。用語は末尾の用語集にもまとめています。
このセクションで分かること:MHA が何をしていて、nn.MultiheadAttention がどんな条件で高速経路を選ぶのか。
MHA(Multi-Head Attention、マルチヘッド注意)は、同じ原稿を $h$ 人の編集者が別々の色眼鏡で読み、編集長が1枚にまとめる仕組みです。記号は次の3つだけ使います。
- $E$(埋め込み次元、整数):1単語あたりの情報の幅。引数名は embed_dim
- $h$(ヘッド数、整数):編集者の人数。引数名は num_heads
- $d$(ヘッド次元、整数):1人が担当する幅で、$d = E / h$
各編集者は、問いのメモ $Q$、見出し札 $K$、中身 $V$(どれも行列、列数は $d$)から次の計算をします。
\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right) V
この式が言っているのは、問いと見出し札の相性を測って赤ペンの濃さに直し、その濃さで中身を混ぜる、ということです。この計算の書き方によって、nn.MultiheadAttention の中には3つの経路があります。
| 経路 | 選ばれる条件 | 中身 |
|---|---|---|
| 手組み | need_weights=True で fastpath の条件を満たさない | 行列積と softmax を順に実行 |
| SDPA 経路 | need_weights=False で fastpath の条件を満たさない | F.scaled_dot_product_attention を呼ぶ |
| fastpath | 推論時の条件をすべて満たす | 専用の _native_multi_head_attention を呼ぶ |
SDPA 経路については、ドキュメントに次の記述があります。
原文:"Set need_weights=False to use the optimized scaled_dot_product_attention"
和訳:need_weights=False にすると、最適化された scaled_dot_product_attention を使う。
出典:PyTorch 2.14 torch.nn.MultiheadAttention(need_weights の説明)
fastpath の条件は、ドキュメントに9つ並んでいます。自己注意であること、batch_first=True、勾配なし、eval モード、kdim と vdim が embed_dim と等しいこと、autocast が無効であること、などです。
ところがインストールされたソース(torch/nn/modules/activation.py)を読むと、ドキュメントにない条件がさらに3つありました。num_heads が偶数であること、float 型のマスクを使っていないこと、bias=False ではないことです。どの経路に入ったかは見た目では分からないので、次の節ではプロファイラで実際に呼ばれた関数を確かめながら測ります。
2. 測る:どの経路が、どれだけ速いか
前節で、呼び方と設定しだいで3つの経路に分かれることを確認しました。経路が分かれるなら、速さも違うはずです。ドキュメントの説明どおり fastpath が最速かどうかを、CPU で確かめます。
このセクションで分かること:6通りの呼び方で実際に通った経路と、CPU での推論時間。
import torch
import torch.nn as nn
from torch.utils.benchmark import Timer
torch.manual_seed(0)
N, L, E = 4, 1024, 448 # 448 は 8 でも 7 でも割り切れる
x = torch.randn(N, L, E)
OPS = ("aten::_native_multi_head_attention", "aten::scaled_dot_product_attention")
def used_op(mha, nw):
with torch.no_grad(), torch.profiler.profile() as prof:
mha(x, x, x, need_weights=nw)
names = {e.key for e in prof.key_averages()}
return next((op.split("::")[1] for op in OPS if op in names), "softmax(手組み)")
def bench(heads, train, nw, fastpath=True):
torch.backends.mha.set_fastpath_enabled(fastpath)
mha = nn.MultiheadAttention(E, heads, batch_first=True).train(train)
ctx = "torch.enable_grad()" if train else "torch.no_grad()"
t = Timer(f"with {ctx}: mha(x, x, x, need_weights={nw})",
globals={"mha": mha, "x": x, "torch": torch})
ms = t.blocked_autorange(min_run_time=3.0).median * 1e3
return ms, used_op(mha, nw)
print(f"torch {torch.__version__}, threads={torch.get_num_threads()}, N={N} L={L} E={E}")
cases = [(8, True, True, True), (8, False, True, True), (8, False, False, True),
(8, False, True, False), (8, False, False, False), (7, False, False, True)]
for heads, train, nw, fp in cases:
ms, op = bench(heads, train, nw, fp)
print(f"heads={heads} train={train!s:5} need_weights={nw!s:5} "
f"fastpath={fp!s:5} {ms:7.1f} ms {op}")
torch.backends.mha.set_fastpath_enabled は、fastpath を全体で有効にするか無効にするかを切り替える関数です。スクリプトを3回実行し、各行の中央値と幅を載せます(forward のみ、N=4, L=1024, E=448)。
| heads | モード | need_weights | fastpath設定 | 実際に通った経路 | 中央値 | 3回の幅 |
|---|---|---|---|---|---|---|
| 8 | train | True | 有効 | 手組み | 331.6 ms | 308〜368 |
| 8 | eval | True | 有効 | _native_multi_head_attention | 269.1 ms | 257〜288 |
| 8 | eval | False | 有効 | _native_multi_head_attention | 238.0 ms | 226〜298 |
| 8 | eval | True | 無効 | 手組み | 354.1 ms | 309〜359 |
| 8 | eval | False | 無効 | scaled_dot_product_attention | 172.7 ms | 165〜177 |
| 7 | eval | False | 有効 | scaled_dot_product_attention | 156.7 ms | 149〜158 |
この環境で最も速かったのは、fastpath を通らずに SDPA 経路へ回った呼び方でした。fastpath(238.0 ms)より SDPA 経路(172.7 ms)のほうが約1.4倍速く、既定の呼び方に近い train の行(331.6 ms)と比べると約1.9倍です。速いはずの経路が遅いと分かったとき、正直に言うと最初は計測ミスを疑って3回測り直しました。幅を見ても、順位は一度も入れ替わっていません。
最後の行は、num_heads=7 という奇数にしただけで、fastpath を有効にしたまま自動的に SDPA 経路へ回ったことを示しています。前節でソースから読んだ偶数条件が、実際に効いている証拠です。
この表は GPU なしの2コア CPU での結果です。GPU 上の挙動や、NestedTensor でパディングを省いたときの fastpath の効果は測っていません。ドキュメントは、NestedTensor を渡すとパディングの割合に比例した高速化が見込めると書いています。自分の環境では、上のスクリプトで両方を測ってから決めてください。
ここまでのまとめ
- nn.MultiheadAttention の中には、手組み、SDPA 経路、fastpath の3経路がある
- fastpath にはドキュメントにない条件があり、num_heads が奇数なら入らない
- 2コア CPU の推論では、SDPA 経路が fastpath より約1.4倍速かった
- どの経路に入ったかは、プロファイラで呼ばれた関数名を見れば分かる
後半は、最速だった SDPA を直接呼ぶ自前実装を組み、本家と同じ答えが出るかを確かめます。
3. SDPAで組み直す:自前のマルチヘッド注意
前節で、最速だったのは nn.MultiheadAttention の外側にある SDPA を呼ぶ経路でした。それなら、最初から SDPA を直接呼ぶ層を自分で書けば、経路の条件分岐に振り回されずに済みます。ドキュメントも nn.MultiheadAttention を「基礎を理解するための参照実装」と位置づけ、新しい構成には部品から組むことを勧めています。
このセクションで分かること:SDPA で組んだ自前実装の中身と、移行で必ず踏む2つの注意点。
1つ目の注意点はマスクの向きです。SDPA のドキュメントにはこう書かれています。
原文:"If migrating from MHA, ensure you invert your boolean mask"
和訳:MHAから移行する場合は、boolマスクを必ず反転すること。
出典:PyTorch 2.14 torch.nn.functional.scaled_dot_product_attention
nn.MultiheadAttention の key_padding_mask は True が無視、SDPA の attn_mask は True が参加です。たとえで言えば、片方は「読むな」の付箋、もう片方は「読め」の付箋です。下の実装では、呼び出し側は MHA と同じ向きのまま渡せるよう、内部で反転しています。
2つ目は dropout です。SDPA は dropout_p を常に適用するので、評価時は呼び出し側で 0.0 を渡すよう、ドキュメントに警告があります。nn.Module の training フラグで切り替えます。
import torch.nn as nn
import torch.nn.functional as F
class SDPAAttention(nn.Module):
def __init__(self, embed_dim, num_heads, num_kv_heads=None, dropout=0.0, bias=True):
super().__init__()
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
self.h, self.h_kv = num_heads, num_kv_heads or num_heads
assert self.h % self.h_kv == 0, "num_heads must be divisible by num_kv_heads"
self.d, self.p = embed_dim // num_heads, dropout
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
self.kv_proj = nn.Linear(embed_dim, 2 * self.h_kv * self.d, bias=bias)
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
def forward(self, x, key_padding_mask=None, is_causal=False):
"""x: (N, L, E)。key_padding_mask は MHA と同じく True = 無視。"""
N, L, _ = x.shape
q = self.q_proj(x).view(N, L, self.h, self.d).transpose(1, 2)
k, v = self.kv_proj(x).view(N, L, 2, self.h_kv, self.d).permute(2, 0, 3, 1, 4)
mask = None if key_padding_mask is None else ~key_padding_mask[:, None, None, :]
o = F.scaled_dot_product_attention(
q, k, v, attn_mask=mask, is_causal=is_causal,
dropout_p=self.p if self.training else 0.0, # 公式Warningの対策
enable_gqa=self.h_kv != self.h)
return self.out_proj(o.transpose(1, 2).reshape(N, L, -1))
num_kv_heads は GQA(Grouped-Query Attention、グループ化クエリ注意)のための引数です。たとえで言えば、8人の編集者がそれぞれ見出し札と中身のコピーを持つのをやめ、2人で1部を共有するやり方です。問いのメモは8人分のままなので読み方の多様さはある程度保ちつつ、コピーの量が減ります。nn.MultiheadAttention にはこの引数がありません。
4. 本家と一致するか、GQAで何が減るか
前節の実装は、マスクの反転と dropout の切り替えを自分で書いた分だけ、間違える余地も増えています。速くても答えが違えば意味がないので、nn.MultiheadAttention の重みをそのまま移植して数値を比べます。そのうえで、GQA にしたとき何が減るかを測ります。
このセクションで分かること:自前実装が本家と一致すること、そして KV ヘッド数を減らすとパラメータと KV キャッシュがどれだけ減るか。
import torch
import torch.nn as nn
from sdpa_attention import SDPAAttention
torch.manual_seed(0)
E, H, N, L = 512, 8, 2, 16
ref = nn.MultiheadAttention(E, H, batch_first=True).eval()
mine = SDPAAttention(E, H).eval()
Wq, Wk, Wv = ref.in_proj_weight.chunk(3)
bq, bk, bv = ref.in_proj_bias.chunk(3)
with torch.no_grad(): # 重みを移植して数値を比べる
mine.q_proj.weight.copy_(Wq), mine.q_proj.bias.copy_(bq)
mine.kv_proj.weight.copy_(torch.cat([Wk, Wv])), mine.kv_proj.bias.copy_(torch.cat([bk, bv]))
mine.out_proj.load_state_dict(ref.out_proj.state_dict())
x = torch.randn(N, L, E)
pad = torch.zeros(N, L, dtype=torch.bool)
pad[0, 12:] = True
y_ref, _ = ref(x, x, x, key_padding_mask=pad, need_weights=False)
diff = (y_ref - mine(x, key_padding_mask=pad)).abs().max().item()
print(f"nn.MHA との最大誤差: {diff:.2e}")
print("\nE=512, num_heads=8 で KV ヘッド数を減らしたとき")
d = E // H
for h_kv in (8, 4, 2, 1):
m = SDPAAttention(E, H, num_kv_heads=h_kv)
params = sum(p.numel() for p in m.parameters())
kv_bytes = 2 * h_kv * d * 2 # K と V、fp16 で 1 トークン 1 層あたり
with torch.no_grad():
shape = tuple(m(torch.randn(N, L, E)).shape)
print(f" kv_heads={h_kv} params={params:,} KVキャッシュ={kv_bytes:5d} B/token/層 出力{shape}")
パディング付きでも、nn.MultiheadAttention との最大誤差は 2.09e-07 でした。float32 の丸め誤差の範囲なので、マスクの向きも含めて本家と同じ計算をしています。
| kv_heads | パラメータ数 | KVキャッシュ(fp16、1トークン1層あたり) |
|---|---|---|
| 8(通常のMHA) | 1,050,624 | 2,048 B |
| 4 | 787,968 | 1,024 B |
| 2 | 656,640 | 512 B |
| 1(MQA) | 590,976 | 256 B |
kv_heads を8から2に減らすと KV キャッシュは4分の1になります。KVキャッシュは、文章を1トークンずつ生成するときに、過去のトークンの Key と Value を保存しておく領域です。長い文脈で VRAM が先に尽きる LLM では、この差が効きます。kv_heads=1 は MQA(Multi-Query Attention)と呼ばれる、全員で1部を回し読みする特別な場合です。
5. 使い続けるか、自前で組むか:判定フロー
前節までで、速さ、自前実装の正しさ、GQA の効果をそれぞれ実測で確かめました。これらを、冒頭の問いに答える形でまとめます。
このセクションで分かること:自分の用途で、nn.MultiheadAttention と SDPA 自前実装のどちらを選ぶか。
nn.MultiheadAttention を使い続けてよいのは、注意の重みを見たいとき、学習や実験の読みやすさを優先するとき、既存の nn.Transformer 系と組み合わせるときです。
自前実装に切り替えるべきなのは、GQA で KV キャッシュを減らしたいとき、RoPE(回転位置埋め込み)のような位置の情報を射影後の Q と K に差し込みたいとき、そして第2節のように SDPA 経路のほうが速いと測れたときです。
6. トラブルシューティング
前節の判定フローで自前実装を選んだ場合に、移行で出会いやすい問題をまとめます。
| 症状 | 原因 | 対処 |
|---|---|---|
| fastpath に入らない | num_heads が奇数、float マスク、bias=False、学習モードなど | ソースの why_not_fast_path の条件を確認する |
| SDPA に移したら結果がずれた | boolマスクの意味が逆 | ~mask で反転する |
| 評価のたびに出力が変わる | SDPA の dropout_p が常に適用される | self.training で 0.0 に切り替える |
| RuntimeError: Explicit attn_mask should not be set when is_causal=True | パディングマスクと is_causal を同時に渡した | 因果マスクとパディングマスクを合成した1つの bool マスクを渡す |
| AssertionError: num_heads must be divisible by num_kv_heads | GQA のヘッド数が割り切れない | num_heads を割り切る num_kv_heads を選ぶ |
| 本家と小数点以下の値がわずかに違う | 選ばれたカーネルで浮動小数点の演算順が違う | 1e-6 程度なら正常。厳密に比べるなら float64 で試す |
4行目は、上の SDPAAttention に key_padding_mask と is_causal=True を同時に渡して実際に出たエラーです。SDPA のドキュメントにも、attn_mask と is_causal の両方を設定するとエラーになると書かれています。
まとめ
nn.MultiheadAttention の中には3つの経路があり、どれに入るかは呼び方と設定で決まります。2コア CPU の推論では、推論用と説明されている fastpath より、need_weights=False で SDPA に回る経路のほうが約1.4倍速いという、説明とは逆の結果になりました。
既製の編集部をそのまま雇うか、自分で編集部を組むか。重みを見たいなら既製品、GQA や速さが要るなら SDPA で自前、迷ったら第2節のスクリプトを自分の環境で回す、というのがこの記事の結論です。「高速」と名の付いた経路でも、測ってみるまでは本当に速いとは限りません。
学習ロードマップ
次に読むなら、SDPA が GPU で使うカーネルの1つである Flash Attention、全体の構造をつかむ Transformer の順がおすすめです。
用語集
| 用語 | 意味 | たとえでの対応 |
|---|---|---|
| MHA | 注意を複数のヘッドに分けて並列に計算し、最後にまとめる仕組み | 複数の編集者と1人の編集長 |
| SDPA | F.scaled_dot_product_attention。最適化済みの注意関数。True が参加 | 「読め」の付箋で動く別の編集部 |
| fastpath | 条件を満たした推論時だけ使われる nn.MultiheadAttention の専用経路 | 条件付きの急ぎ便 |
| GQA | 複数の問いのヘッドで Key と Value を共有する方式 | 2人で1部の見出し札を共有 |
| MQA | 全ヘッドで Key と Value を1組だけ共有する方式 | 全員で1部を回し読み |
| KVキャッシュ | 生成時に過去トークンの Key と Value を保存する領域 | 見出し札のコピーの保管棚 |
計測環境
- Linux x86_64、Intel Xeon 2.80GHz 2コア、メモリ8GB、GPU なし
- Python 3.13.17、PyTorch 2.14.0(CPU 実行、スレッド数2)
- 速度は torch.utils.benchmark.Timer の blocked_autorange(min_run_time=3.0)の中央値を、スクリプト3回分並べたもの
参考文献
- PyTorch 2.14 ドキュメント torch.nn.MultiheadAttention(fastpath の条件、need_weights の説明、参照実装という位置づけ)
- PyTorch 2.14 ドキュメント torch.nn.functional.scaled_dot_product_attention(マスクの向き、dropout の警告、GQA の制約)
- PyTorch チュートリアル(nn.MultiheadAttention のドキュメントが勧める、部品から効率的な層を組む方法)
- Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need", arXiv:1911.02150, 2019(邦題訳:高速なTransformerデコード ── 書き込みヘッドは1つで足りる。MQA を提案)
- Ainslie et al., "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints", arXiv:2305.13245, 2023(邦題訳:GQA ── マルチヘッドのチェックポイントから汎化マルチクエリTransformerを学習する。GQA を提案)
関連記事
前編(公開後にリンクカードに差し替えます)
最後まで読んでいただきありがとうございました。X でも更新情報を流しています。