PyTorchのAttention機構のプロファイリングにおける最新の調査結果が公開された。NVIDIA A100 GPUを用いたトレース解析により、素朴な実装から最適化されたバックエンドまで、各アプローチのGPUカーネル動作と性能特性が詳細に解明された。 まず、基本演算で構成されるNaive Attentionでは、マスク適用時に予期せぬメモリコピーが発生することが確認された。in-place演算への変更によりコピーカーネルを除去可能だが、勾配計算を必要としない推論段階でのみ安全である。一方、PyTorchの標準関数scaled_dot_product_attentionは入力条件に応じて複数のバックエンドにディスパッチする。mathバックエンドは数値安全性を優先するためFP32で動作し、マスク再構築や追加カーネルにより20個の演算を起動する反面、約3.7倍の低速化を引き起こす。 最適化バックエンドでは、カーネル融合による高速化が実現される。efficientバックエンドはxformers由来の単一カーネルでbf16演算を完了させる。FlashAttention-2はオンチップメモリ活用を優先するためプロファイラ上で低占有率と表示されるが、これはグローバルメモリ転送を回避する設計上の必然である。cuDNNバックエンドは問題規模に応じてカーネルを動的生成し、トランスポーズ演算を省略する代わりに、オーバーヘッドがCPUサイドにシフトする特徴がある。 本解析から得られる実践的知見は、プロファイラトレース前に実行予測を立て、実際の結果と乖離する箇所を技術的インサイトとして掘り下げることである。この手法はTransformerや拡散モデルなど大規模モデルの最適化において、隠れたボトルネックの特定とパフォーマンス向上に不可欠な基準となる。