この記事で決められること

読み終えたあと、推論で 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 での推論時間。

exp5_speed.py

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 フラグで切り替えます。

sdpa_attention.py

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 キャッシュがどれだけ減るか。

exp6_gqa.py

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 でも更新情報を流しています。