PyTorch、ROCm FP8 行列乗算の修正をマージ　スケーリングのインデックス誤りと範囲外読み取りを修正

PyTorch は 11 月 6 日、PR #202816 をメインブランチにマージし、ROCm Composable Kernel（CK）バックエンドで行単位の FP8 行列乗算を実行する際のスケーリング設定を修正しました。この変更はコミット記録で確認できますが、現在インストールされている安定版パッケージに修正が含まれているとは限りません。マージ記録

問題は _scaled_mm のスケーリングテンソルにありました。活性値には [M, 1]、重みには [1, N] が使われます。従来のテンプレートはメモリストライドから CK の設定を推定していましたが、連続した [M, 1] テンソルのストライドは (1, 1) になる場合があり、その結果、活性値のスケーリングが列方向に読み取られると誤判定されていました。数値が誤るだけでなく、出力列数 N が行数 M より大きい場合、スケーリングテンソルの割り当て範囲を越えて読み取る可能性もありました。修正では用途に応じて ("Col", "Row") を指定し、単一次元による曖昧さを避けています。修正コミット

追加された回帰テストは、JIT コンパイルと AOTI、バイアスの有無、同一のコンパイル済みモデルで異なる複数の行数を網羅しています。作者によると、MI350X、ROCm 7.0、ck4inductor の環境では、追加された 4 件のテストが修正後にすべて通過しました。修正前は、AOTI のテスト 2 件で要素の不一致率が 98.7% に達していました。これは指定された環境での正しさを示す証拠であり、モデル品質全般や性能の向上を示すものではありません。テストの説明

技術的な背景として、FP8 ではテンソル全体に共通のスケーリングを適用する方法と、行ごとに異なるスケーリング係数を使う方法があります。後者は数値をより細かく制御できますが、カーネルが各スケーリング値を正しく対応付ける必要があります。PyTorch が以前公開した AMD FP8 の技術記事では、こうした量子化の粒度と、関連演算における CK の役割が説明されています。公式の技術解説

運用チームは、デプロイ対象のコミットを確認し、実際に選択されているカーネルを特定したうえで、異なる行列形状を使ってコンパイル前後の出力を比較するとよいでしょう。公式ドキュメントによると、torch.compile の自動チューニングは行列乗算の設定を評価して選択します。そのため、FP8 を有効にしただけでは、ワークロードが必ず影響を受ける経路を通るとは判断できません。コンパイラのドキュメント