TransformersがGQAのデコード最適化をマージ、マスク付き生成時のK/V複製を削減

Transformersは11月4日、PR #49743をメインブランチにマージし、Grouped-Query Attention（GQA）の単一トークンデコードを改善した。この変更は、ネイティブGQAを直接利用できないSDPA経路、特に左側パディングマスクを使うバッチ生成を対象としている。現時点で確認できるのはメインブランチへの取り込みであり、既存のpipインストール版でも効果が得られるとは限らない。マージ記録

GQAでは、複数のクエリヘッドがより少数のキー・バリューヘッドを共有する。しかし従来の経路では、一部の条件下で引き続きrepeat_kvが呼び出され、各層、各デコードステップでK/Vが複製されていた。新しい方法では、同じキー・バリューヘッドを共有するクエリヘッドを、そのヘッドに対応する複数のクエリ行として並べ替える。対応関係を保ったまま出力を元の順序に戻すことで、大きなキー・バリューテンソルの重複した転送を回避できる。プレフィル、複数トークンのクエリ、従来のenable_gqa経路は、これまでの処理を維持する。技術解説

作者はRTX PRO 6000 BlackwellとPyTorch 2.9.1を使用し、Llama-3.1-8Bに形状を合わせたランダム初期化モデルでテストした。バッチサイズ32、プロンプト長4kの場合、1ステップあたりのデコード時間は152.1ミリ秒から54.2ミリ秒に短縮された。これは特定の構成におけるデコード段階の結果であり、リクエスト全体のスループットや、あらゆるモデルでの高速化を保証するものではない。作者はCUDAとCPUのみでテストしたとも明記している。テスト条件と結果

性能はバックエンドの選択にも左右される。PyTorchのドキュメントによると、SDPAは入力に応じて実装を選択し、融合カーネルにはそれぞれ制約がある。また、浮動小数点の結果はカーネルによって異なる場合がある。そのため、エンジニアリングチームがアップグレードを評価する際は、モデル、マスク、バッチサイズ、コンテキスト長を固定し、ステップごとのレイテンシ、メモリ使用量、出力の一貫性を比較する必要がある。PyTorchドキュメント

デプロイ前には、実際にsdpaが使われていることも確認したい。Transformersではattn_implementationで注意機構のバックエンドを指定できるため、今回の変更の効果を他の実行エンジンにそのまま当てはめることはできない。今後は、正式リリースへの収録状況に加え、実際の重み、異なるハードウェア、コンパイルモードで再現できるかを追跡する価値がある。バックエンド設定ドキュメント