TransformersコミュニティがCPUネイティブGQA最適化を提案、局所テストでデコード速度が約2割向上

Transformersコミュニティは10月12日、CPUのAttention経路を最適化する提案を行い、翌日にプロジェクトのコントリビューターが応答した。この提案は use_gqa_in_sdpa の判定を対象としている。提案者によると、現在はAttention Maskがない場合に限り、CPU呼び出しでネイティブのGrouped Query Attention（GQA）が有効になる。パディングを含むバッチや多くの生成リクエストでは、引き続き repeat_kv を通り、K/Vがクエリヘッド数に合わせて複製される。技術提案

この変更の技術的な価値は、Attention計算の前に発生するデータ移動を減らせる点にある。GQAでは複数のクエリヘッドが、より少数のK/Vヘッドを共有する。フレームワークがテンソルを先に展開すると、共有構造を持たせても余分な複製コストが生じる。Hugging Faceのドキュメントによれば、sdpa はPyTorchのAttention実装を使用し、エンジニアは attn_implementation でバックエンドを選択できる。そのため、SDPAを指定するだけでは、内部でこうした複製が省かれるかどうかは判断できない。Attentionバックエンドのドキュメント

提案者はSmolLM2-135Mの形状に合わせたLlama設定でテストした。条件はクエリヘッド9個、K/Vヘッド3個、bf16、バッチサイズ2、左側を2048トークンでパディングしたプレフィックス、CPUスレッド4本。PyTorch 2.14.1では、デコードステップの中央値が71.5ミリ秒から59.3ミリ秒に短縮され、速度は約1.21倍となった。別のマシンでは1.06倍だった。提案者はlogitsとgreedy decodingで生成されたトークンが一致し、関連テストも通過したと報告しているが、いずれも提案者による局所的な検証である。テスト条件と結果

提案では、PyTorch 2.9以降のCPUカーネルはマスク付きのネイティブGQAを処理できるとし、K/Vに勾配が不要であること、コンパイルまたはトレースされていないこと、K/Vヘッドの次元が一致することを適用条件としている。コントリビューターは、コンパイルとトレースを除外する理由を尋ね、継続的な蓄積条件の保守性にも疑問を示した。現在もIssueはオープンであり、正式リリースでの改善とはまだ言えない。エンジニアは互換性テストとマージ状況を追い、自身のCPU、シーケンス長、バッチ設定で測定する必要がある。また、ここでのデコード結果から生成全体のレイテンシを直接推定することはできない。コントリビューターの応答