PyTorch Profiling 完全解説（第3回）：Attentionの実装別カーネル挙動をプロファイルで読む

記事のサマリー（TL;DR）

masked_fill を masked_fill_（in-place）に1行変更するだけでMemcpyカーネルを1本削減
SDPA mathバックエンドはFP32 CUDA coreで動作し、naive実装比3.7倍遅い；FlashAttention-2は13%occupancyでも最速
cuDNNバックエンドはカーネル自体はfuseされるが、plan選択コストがCPU側に214µsとして現れる

生成AIモデル・推論基盤を持つ日本企業がAttentionプロファイリングで得られる実践的知見

LLMの推論・ファインチューニング基盤を社内に持つ企業や、Hugging Face / PyTorchベースのモデルサービングを運用する開発チームにとって、この記事の知見は直接コストに響きます。Attentionはトランスフォーマー系モデルの全レイヤーで繰り返されるため、1フォワードパスあたり1カーネル削減できれば、レイヤー数×バッチ数のスケールでレイテンシが積み上がります。

日本国内でも、kintone・Salesforce・SmartHR等のSaaS上に構築した業務データをベクトル検索・RAGで活用するケースが増えており、そのバックエンドにはAttentionベースのエンコーダが使われます。推論コストを下げるには「どのSDPAバックエンドを選ぶか」よりも「なぜそのバックエンドがそう動くか」を理解することが先決で、本記事はそのための読み方を体系的に示しています。

また、torch.no_grad() 下でのin-place操作の安全性や、occupancyの数値だけでパフォーマンスを判断しないという知見は、GPUインフラ（AWS / GCP）上でのモデル最適化に携わるエンジニアが実際の意思決定で使える具体的な指針です。

詳細

Naive Attention：プリミティブ操作の集合体

Attention の計算は以下のステップで構成されます。

スコア行列の構築：matmul(q, k.T)
スケーリング：scores * scale
Causal maskの適用：scores.masked_fill(mask, "-inf")
Softmaxで正規化：softmax(scores) → attention weights
値の再重み付け：matmul(attn, v)

これをそのままPyTorchで実装した NaiveCausalAttention をプロファイルすると、CPU レーンには予想通りの操作列（matmul → mul → masked_fill → softmax → matmul）が現れます。

ところがGPUレーンを展開すると、6本のカーネルが確認できます。

#	カーネル	対応する操作
1	matmul	q × k.T
2	mul	スケーリング
3	Memcpy	🤔（予期しない）
4	masked_fill	Causal masking
5	softmax	Softmax
6	matmul	attn × v

5本は想定通りですが、Memcpy が1本余計に発生しています。原因は masked_fill が out-of-place（デフォルト）で動作し、テンソルのコピーを作ってから処理を適用するためです。

In-place操作でMemcpyを1本削減

masked_fill を masked_fill_（末尾のアンダースコアがPyTorchのin-place規約）に変えるだけで、GPUレーンからMemcpyカーネルが消えます。変更はこの1行だけです。

- scores = scores.masked_fill(mask, float("-inf"))
+ scores.masked_fill_(mask, float("-inf"))

CPUレーンでもmaskingステップ内のop数が大幅に減少します。

なぜ普段はout-of-placeがデフォルトなのか
autograd（自動微分）は逆伝播のためにフォワードパスのテンソル値を記憶します。in-place操作でその値が上書きされると逆伝播が壊れます。今回は torch.no_grad() 配下なので安全です。in-place操作はメモリコピーを行わないため、時間だけでなくVRAM消費も削減できます。logitsのような大きなテンソルには特に効果的です。

Scaled Dot Product Attention（SDPA）：1行で全部まとめる

PyTorchには公式の実装があります。

from torch.nn import functional as F
F.scaled_dot_product_attention(q, k, v, is_causal=True)

SDPA は内部で以下の4つのバックエンドにディスパッチします。

from torch.nn.attention import SDPBackend

BACKENDS = {
    "math":      SDPBackend.MATH,
    "flash":     SDPBackend.FLASH_ATTENTION,
    "efficient": SDPBackend.EFFICIENT_ATTENTION,
    "cudnn":     SDPBackend.CUDNN_ATTENTION,
}

torch.nn.attention.sdpa_kernel コンテキストマネージャーで特定バックエンドに固定し、個別プロファイルが可能です。

SDPA mathバックエンド：正確だが3.7倍遅い

プロファイル結果

指標	Naive in-place	SDPA math
CUDA time avg（`*_fwd`）	1.955 ms	7.239 ms
Self CUDA time total	7.194 ms	27.279 ms

1行で書けるのに3.7倍遅い理由は3つあります。

① Tensor Coreが使われない
Naive実装のカーネル名に含まれる s16816 はbfloat16 Tensor Core matmul（16×8×16命令）のシグネチャです。一方、mathバックエンドのカーネル名は sgemm、つまりFP32のCUDA coreによるmatmulです。mathバックエンドは精度確保のためbf16をFP32にアップキャストし（データ量が倍増）、Tensor Coreを使いません。

② Causal maskを毎回再構築
is_causal=True を指定しても、mathバックエンドはすべてのフォワードパスで以下の処理を実行します。

aten::ones → aten::tril：下三角行列を構築
aten::scalar_tensor → aten::fill_：-inf 値を生成
aten::where：加算バイアス（0 or -inf）に変換

Naive実装ではmaskを事前生成して使い回していましたが、convenientなフラグを使うと逆に毎回再構築されます。

③ safe softmax
通常の aten::softmax ではなく aten::_safe_softmax が呼ばれ、全エントリが -inf の行（完全マスク）でNaNが生じないよう追加カーネルが走ります。

総カーネル数はNaive（5本）に対してmathバックエンドは20本。mathバックエンドは「常に正しく動く」リファレンス実装であり、速さではなく安全性が目的です。

SDPA efficientバックエンド：1本のfusedカーネル

fmha_cutlassF_bf16_aligned_64x64_rf_sm80

カーネル名を分解すると以下の通りです。

パーツ	意味
`fmha`	Fused Multi-Head Attention
`cutlassF`	NVIDIAのCUTLASS（Tensor Core GEMM テンプレート集）使用、Fはforward
`bf16_aligned`	bfloat16で動作（FP32アップキャストなし）
`64x64`	タイルサイズ
`rf`	作業データをregister file（最速メモリ）に保持
`sm80`	Ampere世代（A100のcompute capability 8.0）向けコンパイル

20カーネルが1本に集約されます。これはMeta xformers ライブラリ由来のメモリ効率的Attentionカーネルです。

SDPA flashバックエンド：なぜFlashAttentionが必要か

HBMトラフィックがAttentionのボトルネック

mathバックエンドの本当の問題はカーネル数ではなく、カーネル間のデータ受け渡しにあります。

q × k.T で生成されるスコア行列のサイズは [seq, seq]
シーケンス長4096の場合、1ヘッドで 4,096 × 4,096 ≈ 1,600万要素
この行列をHBM（GPUのメインメモリ）に書き出し、スケーリングで読み直し、maskで再書き込み、softmaxで再読み込み、……という往復が続く

Attentionのコストはmatmul自体ではなくHBMへの往復トラフィックが支配しています。

FlashAttentionのアイデア

FlashAttentionはk, vをタイル単位で処理しながら「online softmax」でrunning sumを保持し、出力を1タイルずつ積み上げます。[seq, seq] のスコア行列はHBMに書かれず、オンチップのみで完結します。これにより全Attention処理が1本のfusedカーネル（pytorch_flash：FlashAttention-2実装）に収まります。

occupancy 13%でも最速な理由

Flash kernelはプロファイラ上で13%のoccupancyと表示されます。これは通常「SM利用率が低い＝遅い」と解釈されますが、ここでは逆です。

各ブロックが255レジスタ × 128スレッド = 32,640レジスタを使用
Ampere SMの最大レジスタ数（65,536）に対して2ブロックしか同時ロードできない
4 warps × 2 blocks = 8 resident warps（最大64の13%）

この「重さ」は意図的なものです。高occupancyはレイテンシを隠蔽しますが、それは各操作そのものを効率化しません。Flashはレジスタと共有メモリをattentionタイルのオンチップ保持に使い、HBMトラフィックを根本から排除しています。occupancyの数値だけでパフォーマンスを判断してはいけない典型例です。

cuDNNバックエンド：生成されるカーネルとCPUコスト

cuDNNバックエンドも1本のfused kernelを生成します。カーネル名の例：

cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0

パーツ	意味
`cudnn_generated`	事前コンパイル済みバイナリではなく、問題に合わせて生成
`flash_fprop`	FlashAttentionスタイルのフォワードパス
`wmma_f16`	WMMA（warp-level MMMA）API、16bit Tensor Core経路
`knob_6`	cuDNNが選択した事前チューニング済みコンフィグ
`128x64x64`	選択されたタイルサイズ

Flash/efficientとの技術的差異

transposeなし：Flash/efficientはテンソル整形のために4本のメタデータtransposeを挿入しますが、cuDNNは [B, H, S, D] レイアウトを直接処理するカーネルを生成します
cuLaunchKernelEx使用：他の全カーネルが cudaLaunchKernel（ランタイムAPI）を使うのに対し、cuDNNはドライバーレベルの拡張APIを使用。これによりCUPTIがoccupancyを測定できず、プロファイラ上は0%と表示されます（実際のoccupancyは240レジスタ × 256スレッド ≈ 12.5%でFlashと同等）

CPUコストの比較

バックエンド	CUDA avg時間	CPU avg時間
efficient	277.9 µs	117 µs
flash	146.8 µs	138 µs
cudnn	186.3 µs	214 µs

cuDNNはGPU時間ではflashとefficientの中間ですが、CPU時間は最も大きい（214 µs）。transposeがゼロなのに、「knob」選択・プラン準備のためのランタイムエンジンコストがCPU側に積み上がるためです。ATen opが少ないからといってCPU負荷が少ないわけではなく、仕事がライブラリ内部の不透明なバーとして隠れているだけです。

全バリアント比較まとめ

バリアント	変更点	カーネル数/forward	トレースが示した知見
Naive attention	プリミティブから手書き実装	6	out-of-place `masked_fill` による隠れたMemcpy
Naive in-place	`masked_fill_` に変更	5	1行でMemcpyカーネルを削除
SDPA math	mathバックエンド固定	20	FP32 CUDA core、毎回mask再構築、`_safe_softmax`。正確だが3.7倍遅い
SDPA efficient	efficientバックエンド	1	`fmha_cutlassF` 1本に集約、bf16 Tensor Core
SDPA flash	flashバックエンド	1	`pytorch_flash`（FlashAttention-2）。13%occupancyでも最速
SDPA cuDNN	cuDNNバックエンド	1	問題特化型生成カーネル。transposeなし、CPUコスト214µs

シリーズを通じて得られる習慣

この3回シリーズを通じて強調されてきたのは、「まず予測し、それからトレースを見る」という習慣です。

期待するカーネル構成を言語化してからトレースを開く
予想と一致しない点こそ最も重要な情報
隠れたMemcpy、addmm epilogue、20本のmathバックエンド、Flashの「おかしな」occupancy、cuDNNの太いCPUバーはすべて、予測の不一致から発見されたもの

プロファイリングはGPUエキスパートのための特殊技術ではありません。「なぜこうなっているのか」と問い続ける習慣です。