PyTorchアテンションのプロファイリング：Flash・cuDNN・Mathバックエンド徹底比較

はじめに：アテンションをプロファイラで覗く

Transformerベースのモデルを扱っていると、アテンション演算の性能がモデル全体の速度を左右することを実感します。しかし、いざプロファイラを開くと、無数のカーネルと複雑なタイムラインに圧倒されがちです。本記事はHugging Faceブログの「Profiling in PyTorch」シリーズ第3弾として、アテンションを自作し、PyTorchのSDPA（Scaled Dot Product Attention）バックエンドと比較しながらプロファイラトレースの読み方を解説します。

シリーズの目標はシンプルです。プロファイラトレースを「読める」開発者になること。Part 1では基本演算を、Part 2では線形レイヤーとMLPをプロファイリングしました。今回はアテンションです。各バックエンドがプロファイラでどう異なって見えるか、実際に確認してみましょう。

参考: 本記事の実験はすべてNVIDIA A100-SXM4-80GB GPUで実施されており、元のスクリプトはHugging Faceブログで確認できます。

1. 素朴なアテンション実装とプロファイリング

まずアテンションを基本演算で直接実装してみます。アテンションは以下のステップで構成されます：

1. スコア計算: matmul(q, k.T)
2. スケーリング: scores * scale
3. 因果マスク適用: scores.masked_fill(mask, "-inf")
4. ソフトマックス: softmax(scores)
5. 値の再重み付け: matmul(attn, v)

import torch
import torch.nn as nn
import math

class NaiveCausalAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.scale = 1.0 / math.sqrt(head_dim)

    def forward(self, q, k, v, mask):
        # q, k, v: [batch, heads, seq, head_dim]
        scores = torch.matmul(q, k.transpose(-2, -1))  # [batch, heads, seq, seq]
        scores = scores * self.scale
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)  # [batch, heads, seq, head_dim]
        return out

このモジュールをプロファイリングすると、CPUレーンにmatmul、mul、masked_fill、softmax、matmulが順に現れます。GPUレーンを展開すると、予期しないカーネルが1つ見えます。それが**メモリコピー（Memcpy）**カーネルです。

隠れたMemcpyの正体

masked_fillはout-of-place演算です。PyTorchはデフォルトでテンソルをコピーしてから演算を適用し、結果を返します。この過程でMemcpyカーネルが発生するのです。これをin-place演算（masked_fill_）に変えるとどうなるでしょうか？

def forward(self, q, k, v, mask):
    scores = torch.matmul(q, k.transpose(-2, -1))
    scores = torch.mul(scores, self.scale)
    scores.masked_fill_(mask, float("-inf"))  # in-placeに変更
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v)
    return out

GPUレーンからMemcpyカーネルが消えたことを確認できます。たった1行の変更でフォワードパスごとにカーネルを1つ削減したわけです。レイヤーが多くシーケンスが長いLLMでは、この差が積み重なって大きな性能向上につながります。

注意: in-place演算はautogradが逆伝播のために保存したテンソル値を上書きする可能性があります。torch.no_grad()コンテキスト内でのみ安全に使用してください。

2. SDPAバックエンド別プロファイラトレース比較

PyTorchはアテンションパイプライン全体をF.scaled_dot_product_attentionの1行で提供します。この関数は内部的に複数のバックエンドから最適なものを選択しますが、torch.nn.attention.sdpa_kernelコンテキストマネージャで特定のバックエンドを強制できます。

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

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

各バックエンドのプロファイリング結果を比較すると以下の通りです：

mathバックエンド：リファレンス実装の罠

最初に確認すべきはmathバックエンドです。1行の関数に置き換えたのだから速くなるはず、と思いきや、実際には自作実装より3.7倍遅いのです。理由は明確です：

• Tensor Core未使用: sgemmカーネルはFP32 CUDAコアで実行されます。bf16入力もFP32にアップキャストされ、データ移動量が2倍になります。
• マスクを毎回再生成: is_causal=Trueフラグは利便性を提供するだけで、内部的には呼び出しごとに[seq, seq]三角行列を生成します。
• _safe_softmax: 完全にマスクされた行でNaNを防ぐ安全装置が追加カーネルを発生させます。

mathバックエンドは正確性が保証されたリファレンス実装です。速くはありませんが常に動作します。他のバックエンドはこの20個のカーネルを1つの融合カーネルに畳み込むことを目指しています。

flashバックエンド：低占有率が罠でない理由

FlashAttention-2は最速のバックエンドですが、プロファイラは13%の低い占有率（occupancy）を報告します。これは最適化不足ではなく意図された設計です。

Flashはタイル単位でkとvを走査し、オンラインソフトマックス手法で全体の[seq, seq]行列をHBMに書き込みません。代わりに多くのレジスタと共有メモリを使い、データをオンチップに保持します。ブロックあたり128スレッド × 255レジスタ = 32,640レジスタが必要で、65,536レジスタを持つSMには2ブロックしか常駐できません。結果として8ワープのみがアクティブになり、約13%の占有率になります。

占有率が低いからといって性能が悪いわけではありません。Flashはレジスタと共有メモリを意図的に多く使い、データ再利用を最大化し、グローバルメモリアクセスを最小化しています。

cuDNNバックエンド：生成カーネルの光と影

cuDNNはNVIDIAのディープラーニングライブラリで、問題ごとにカーネルを生成・チューニングします。カーネル名cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0から分かる通り、事前コンパイルされたバイナリではありません。

利点: ネイティブの[B, H, S, D]レイアウトを直接消費するため、aten::transposeが一切発生しません。flashとefficientはそれぞれ4つのメタデータ転置を挿入します。

欠点: CPUオーバーヘッドが大きいです。aten::scaled_dot_product_attentionのself timeが全体実行の26%を占め、これはcuDNNランタイムが呼び出しごとに「knob」検索を実行するためです。プロファイラトレースがきれいに見えても、作業が消えたわけではありません。作業がプロファイラで分解できない場所に移動したのです。

実務のヒント: アテンション性能最適化ではまずF.scaled_dot_product_attentionを使いましょう。特定のバックエンドがより良い性能を示すなら、sdpa_kernelコンテキストで固定できます。ただし必ずプロファイラで実際のボトルネックを確認してから決定してください。

まとめ：推測し、確認し、不一致を見つけよ

このシリーズ全体を貫く1つの習慣があります。まず推測し、それから確認することです。

プロファイラトレースを開く前に「これらのカーネルが見えるはずだ」と声に出してみましょう。そしてトレースを開いて予想と違えば、その不一致を画面で最も興味深いものとして扱ってください。本記事のすべてのインサイト — 隠れたMemcpy、mathバックエンドの20カーネル、flashの低占有率、cuDNNのCPUオーバーヘッド — はすべて推測とトレースの不一致から生まれました。

プロファイリングはGPU専門家だけの技術ではありません。「ちょっと待って、なぜこうなる？」と問いかけ、答えがしっくりくるまで覗き込む習慣にすぎません。これで皆さんは自分のモデルでその習慣を実践する語彙と反射神経を身につけました。トレースを開き、推測し、不一致を探しに行きましょう。

日本の開発エコシステムにおける適用文脈

日本ではLLMやビジョンモデルのファインチューニングやサービングを行うチームにとって、アテンション最適化はコスト削減に直結します。特にA100/H100 GPUをクラウドで借りる環境では、F.scaled_dot_product_attentionへの移行が即座の性能向上をもたらす可能性があります。ただしcuDNNバックエンドはCPUオーバーヘッドが大きいため、バッチサイズが小さい、あるいは呼び出し頻度が高いサービング環境ではflashバックエンドの方が有利な場合があります。実際のサービスに適用する前に、必ず自社のワークロードでプロファイリングしてください。

あわせて読みたい

• OracleからPostgreSQLへ、AIがマイグレーションを革新する方法（feat. Azure）
• CSSハイライト疑似要素完全ガイド search-textからtarget-textまで