PyTorchの注意機構をプロファイラで解剖

Hugging Faceは2026年7月10日、PyTorchのプロファイリング手法を解説する連載の第3回を公開しました。今回のテーマは注意機構（Attention）で、手書きの素朴な実装からSDPA（Scaled Dot Product Attention）の各バックエンドまでを、プロファイラのトレースで読み解きます。狙いは高速化テクニックの網羅ではなく、各手法がプロファイラ上でどう違って見えるかを掴むことにあります。

素朴な実装では、行列積・スケーリング・マスク・softmaxといった基本演算に分解して注意を計算します。トレースを開くと想定外のメモリコピーが1つ紛れ込んでおり、その原因は非破壊的なmasked_fillにありました。これをインプレース版のmasked_fill_に置き換えるだけでフォワードごとにカーネルが1つ減り、大規模モデルでは層の数だけ効果が積み上がります。

PyTorchはこの一連の処理をF.scaled_dot_product_attentionという1行に集約しています。ところがmathバックエンドに固定すると、カーネル数は5から20へ増え、約3.7倍も遅くなりました。FP32へ格上げしてTensor Coreを使わず、呼び出しのたびに因果マスクを作り直し、NaN対策の安全なsoftmaxを走らせるためです。

一方、efficient・flash・cuDNNの各バックエンドは、20個のカーネルを実質1つの融合カーネルにまとめます。とくにflashはFlashAttention-2の実装で、巨大なスコア行列をメモリに書き出さず、オンラインsoftmaxでチップ上だけで処理を進めます。これがHBMへの往復を避け、bf16のままTensor Coreで走らせる鍵になります。

flashはプロファイラ上で占有率が約13%と低く表示され、一見すると効率が悪そうに見えます。しかしこれはレジスタと共有メモリを意図的に多く使い、データをチップ上に留める設計の結果であり、最速という事実は変わりません。筆者は連載を通じ、まず結果を予想してからトレースを開き、食い違いこそ最も面白い手がかりとして掘り下げる習慣を勧めています。