PyTorchコミュニティが融合処理の修正を提案、量子化LLMの単一トークンデコードが大きな中間バッファーで遅くなる可能性

PyTorchコミュニティは10月10日、Inductorのパフォーマンス問題を公開した。max-autotune関連の設定を使用すると、前段の演算子を融合するかどうかの判定によって、単一行の行列乗算で、本来は省略できる大きなFP32中間バッファーが生成されることがある。報告者はRTX 5090で2.14.1と2.16開発版の両方において、特定の演算子の組み合わせで速度が低下することを確認し、最小再現コードと生成コードを添えた。問題報告

報告の分析によると、行列乗算のMが1に固定されている場合、この経路では演算がブロードキャスト乗算と加算に分割される。両側の入力でそれぞれの演算数がしきい値以下でも、融合後にGPUのデフォルト値である30を超えると、コンパイラーが完全な(1, K, N)積を実体化し、その後、リダクションカーネルに読み込ませることがある。K=9216、N=2560の場合、約94 MBのデータが追加で書き込まれ、読み戻される。演算子の数が少し増減すると、かえってこの条件を回避できる場合もある。再現手順と分析

公式ドキュメントによると、max-autotuneはTritonまたはテンプレートベースの行列乗算を使用し、GPUではCUDA Graphsがデフォルトで有効になる。一方、max-autotune-no-cudagraphsではグラフのリプレイが無効になる。この違いを利用すれば、カーネル生成とリプレイの効果を分けて比較できる。今回の事例からは、量子化デコードを評価する際に、活性化関数、逆量子化、行列乗算をまとめて計測すべきだと考えられる。乗算だけを個別にテストすると、融合によるメモリ転送量の増加を見落とす可能性がある。公式ドキュメント

ドラフトPR #200540は、完全な積をメモリに書き出さずに済むよう、より小さなブロードキャスト入力を先に実体化することを提案している。作者が実施したRTX 4090でのテストでは、SwiGLUの前段演算を伴うint4重み量子化のケースで、実行時間が179.5マイクロ秒から27.1マイクロ秒に短縮した。また、バッファー割り当てと数値比較のテストも追加されている。ただし、これは個別カーネルの計測であり、モデル全体のトークン処理量の向上にそのまま換算することはできない。修正提案

確認時点で、この提案はドラフトのままでレビューはなく、一部のCIは承認待ちだ。エンジニアはまず、生成コードで上記の積のバッファーが割り当てられているかを確認し、コンパイルモード、レイテンシ、数値結果を比較できる。今後は、上流でのレビュー状況と、正式版に修正が取り込まれるかを追う必要がある。提案の状況