Hugging Face、A100でのPyTorch attentionプロファイリング：naiveカーネルからflashまで

Hugging Faceの「Profiling in PyTorch」シリーズ第3回のポストでは、torch.profilerのトレースを読み解き、A100上でのattention実装を比較している。対象となるのは、手作業で構築したnaiveなバージョン、in-placeのバリアント、および4つのSDPAバックエンドだ。ここから得られる知見は実践的である。開発者はトレース内の隠れたカーネル（out-of-placeなmasked_fillがforward passごとにメモリのコピーを暗黙的に追加している点など）を発見し、最速のパスを見極めることができる。各バックエンドの実行時間は、flashの146.8µsに対し、efficientは277.9µs、cuDNNは186.3µsとなっている。

出典： huggingface.co ↗

たった一行の変更で、各フォワードパスから一つのカーネルを丸ごと削り取った。

Hugging Face

なぜ重要か

→ masked_fill内の隠れたメモリーコピーが、層ごとのattention性能を静かに劣化させる。
→ SDPAバックエンドの選択は、A100全体でカーネル数とレイテンシーを劇的に変化させる。
→ プロファイラートレースは、生のコード検査では見えない最適化の成果を明らかにするものだ。

attentionカーネルのプロファイリング