PyTorch、Tritonのデコード性能低下問題をクローズ　3.9ブランチへの修正のリバートで特定のattentionカーネルが高速化

PyTorchのTritonデコード性能低下の追跡は10月26日に収束し、メンテナーが大量に出ていた警告メッセージも解消したと確認して、問題をクローズした。これは先行して行われたコンパイラー変更のリバートと検証の続報であり、当日に新たな不具合が見つかったわけではない。推論環境のアップグレードを準備しているチームにとって重要なのは、実際にインストールされているビルドを確認することだ。問題の記録

問題は当初、SGLangのワークロードで明らかになり、報告者はモデルの重みを必要としない独立した再現コードも提供した。Tritonの修正説明によると、マスク付き非同期コピーのpredicateオペランドの変更が、別のグローバルロードの変更と相互作用し、バックエンドの命令スケジューリングに影響した。GB200上の指定されたGQAデコードカーネルでは、レイテンシが約10.9マイクロ秒から17.4マイクロ秒に増加した。レジスター使用量、spill、数値結果に変化がなかったことから、リソース使用量だけでは性能低下を説明できないことが示された。修正の分析

対応として、Tritonのrelease/3.9.xで関連変更をリバートし、修正は10月20日にマージされた。その代わり、別のパイプライン化コピーのチュートリアル例で得られていた効果は失われた。これは、コンパイラー最適化の結果がカーネルの形状やスケジューリングによって異なり得ることを示している。Triton公式のattentionチュートリアルには、実行可能なカーネルと、精度や形状を変えた測定例も掲載されており、こうした性能差を理解する際の参考になる。マージ記録、公式チュートリアル

その後のGB300での比較では、指定ケースのレイテンシが、以前のビルドでの16.86マイクロ秒から10.72マイクロ秒に回復した。別の参加者も、先に報告されていたSGLangワークロードの性能がGB200とGB300で回復したと確認している。ただし、公開された詳細な数値の中心は合成カーネルのもので、完全なサービスのスループット向上に換算することはできない。メンテナーは、環境に3.9.0という同じバージョン表記の古いwheelが残っていたため、修正の効果を測定できなかったこともあった。エンジニアはパッケージの取得元とビルドを確認したうえで、実際のKV設定、バッチサイズ、エンドツーエンドのレイテンシで再測定するとよい。検証と制約