本記事では、プロファイラトレースを使用してPyTorchにおけるアテンションメカニズムをプロファイルし、パフォーマンスのボトルネックと最適化の機会を特定する方法を示します。カーネルの動作を説明するために、単純なアテンション実装とPyTorchの組み込みScaled Dot Product Attention (SDPA)を比較しています。

  • 単純な因果的アテンションには、matmul、スケーリング、マスク処理、softmax、およびアウトオブプレース操作による予期せぬメモリコピーを含む複数のカーネルが含まれます。
  • masked_fillをインプレースのmasked_fill_に置き換えると、GPUトレースから不要なメモリコピーカーネルが削除されます。
  • PyTorchのSDPA関数は、入力タイプとハードウェアに基づいて、Flash AttentionやcuDNNなどの最も高速な利用可能なバックエンドに自動的にディスパッチします。
  • 個々のSDPAバックエンドをプロファイルすると、固有のカーネルシグネチャが明らかになり、開発者はどの実装がアクティブかを確認できます。

これらのプロファイラトレースを理解することで、ユーザーは冗長なメモリ操作を排除し、PyTorchの効率的なアテンションバックエンドを活用してトランスフォーマーモデルを最適化できます。