PyTorch、CUDAのインデックス書き込みエラーを修正　特定の連続テンソルレイアウトで範囲外書き込みの可能性

PyTorchで当初は注意機構のコンパイル異常に見えた問題が、11月7日、CUDAのindex_put_書き込みカーネルにあると特定されました。11月5日の報告では、Gemma 4で静的キャッシュ、torch.compile、パック形式のGrouped-Query Attention（GQA）を組み合わせた際に、誤った出力が生じる可能性が指摘されました。その後、メンテナーがモデルやコンパイラーに依存しない最小再現例を示し、根本原因を低レベルのテンソル操作に絞り込みました。公式の問題記録

長さが1の次元には特殊なストライドを設定でき、そのテンソルも連続レイアウトとして扱われる点が鍵です。影響を受けたソート書き込みカーネルは、その次元のストライドを使ってスライス間を移動していたため、テンソルの範囲外にデータを書き込み、本来更新されるはずの位置が更新されませんでした。メンテナーが示した再現例では、期待値が[0,1,0,1]であるのに対し、実際の値は[0,1,0,0]となり、別の値は基底ストレージ内の誤った位置に書き込まれていました。根本原因と再現例

発生条件から、問題の追跡が難しかった理由も分かります。モデルテストで決定論的アルゴリズムが有効になっていたため、キャッシュ更新が該当カーネルを通りました。公式ドキュメントによると、torch.use_deterministic_algorithms(True)を設定すると、対応する操作は決定論的な実装に切り替わり、Inductorの動作も調整されます。そのため、出力を比較する際は、この設定を実行環境の条件に含める必要があります。決定論的アルゴリズムのドキュメント

メンテナーによると、特殊なレイアウトでの加算書き込みと、決定論的モードでの複数要素の上書きが、いずれも影響を受けた経路を通る可能性があります。範囲外書き込みは、不正なメモリアクセスを引き起こしたり、隣接するテンソルを破損させたりするおそれがあります。問題記録は11月7日にクローズされ、修正コミット1e69880へのリンクが追加されました。ただし、これによって既存のインストール環境にも修正が適用されたとは限りません。修正状況

推論サービスを運用するチームは、修正を含むパッケージやバックポート版の提供状況を追い、最小再現例とスライディングウィンドウキャッシュのテストをアップグレード検証に加えるべきです。今回の事例からは、通常実行とコンパイル実行の出力比較に加え、異なるストライドや決定論的設定もテスト対象に含める必要があると考えられます。現時点の公開情報だけでは、影響を受ける全バージョンやモデルを特定できません。