Codexによる自動リサーチ：232倍高速なQR分解カーネルの実現

エグゼクティブ・サマリー

OpenAIのCodexを使用した自動リサーチ・ループを実装することで、開発者はバッチ処理された正方コンパクト・ハウスホルダーQR分解において、torch.geqrfのベースラインに対して232倍の高速化を達成しました。この成功は、汎用的なライブラリ呼び出しからブロック化されたハウスホルダー・アルゴリズムへの移行、局所解を避けるための「候補のビーム（beam of candidates）」の活用、そしてGPU Mode popcorn CLIとModalによるタイトなフィードバック・ループの活用によってもたらされました。

最適化の課題：QR分解

目標は、FP32 CUDA行列に対するバッチ処理された正方コンパクト・ハウスホルダーQR分解を実装することでした。出力には、H行列（上三角部分がRで、下三角部分にハウスホルダー・ベクトルを格納）と、反射係数のtauベクトルからなるコンパクトな表現が必要でした。

直列的なボトルネック

標準的なハウスホルダーQRは本質的に直列的です。反射器 $j+1$ は、反射器 $j$ によって生成された行列に依存します。この逐次的な依存関係は、Tensor Coreの有効な利用を妨げます。なぜなら、計算が「行列-行列」形式ではなく「行列-ベクトル」形式にとどまるため、GPUの最も強力な計算ユニットがアイドル状態になってしまうからです。

ブロック化されたハウスホルダー・ソリューション

直列的なボトルネックを克服するために、ブロック化されたハウスホルダー・アルゴリズムが採用されました。このアプローチでは、直列的な作業を幅 $b$ の狭いパネルに限定します。その後、$b$ 個の反射器は単一のランク-$b$ 更新（WY表現）に圧縮され、行列の残りのブロックを3つの連続するGEMM（General Matrix Multiplications）を使用して更新することが可能になります。これにより、計算の大部分がTensor Coreで最大限の効率で実行可能な形式へと変換されます。

自動リサーチ手法：「ループ・エンジニアリング」

開発者は高頻度の反復ループを利用し、14日間で1,500回以上の提出を行いました。このプロセスは、LLM（CodexとClaude）と構造化されたハーネスの組み合わせに依存していました。

ハーネスの設定と制御

ツール類: 開発者はベンチマークとリーダーボードへの提出のためにpopcorn CLIを使用し、プロファイリングとnsys/NCU分析のためにModalを使用しました。
ロギング: log.md ファイルによって、すべての提出、そのステータス（accept/reject）、および形状ごとのタイミングを追跡し、冗長な実験を防ぎました。
目標指向のプロンプティング: Codexの/goalコマンドを使用して、定量的ターゲット（例：「現在の最高記録のn=512のタイミングを上回る」）を設定することで、モデルが数時間または数日間、自律的に反復を行うことを可能にしました。
監視: /btw および /side コマンドにより、開発者はメインの最適化ループを停止させることなく、エージェントの進捗や現在の仮説を問い合わせることができました。

ビーム・サーチによる局所解の回避

パフォーマンスが3,000 $\mu$sの境界に達すると、モデルは頻繁に局所解に陥り、構造的な革新ではなく、些細なパラメータ調整に終始するようになりました。これを解決するために、開発者は**候補のビーム（beam of candidates）**を実装しました。

単一の現行記録保持者ではなく、エージェントは3〜5個の「アイデア・ファミリー」を維持しました。
これにより、高リスクな構造的変更が、初期段階で現在のベストよりもパフォーマンスが低かったという理由だけで、早期に破棄されることを防ぎました。
ビームには通常、1つの「exploit（活用）」ビーム（微調整）、1つの「near-miss（惜しい）」ビーム、および1つの「structural/high-risk（構造的/高リスク）」ビームが含まれていました。

カーネルの技術的進化

カーネルは、最終的な1,805 $\mu$sの幾何平均に到達するために、10個の主要な構造的ブレイクスルーを経て進化しました。

Stage	Change	Impact
1	`torch.geqrf`	ベースラインの開始点 (>108.8k $\mu$s)
2	Blocked WY QR (n=512)	パネル・ファクターライゼーションと残りの更新を導入
3	Blocked route (all shapes)	すべての行列サイズに対してブロック化ロジックを適用
4	Triton panels	カスタムのpanel16/32カーネルを実装
5	Cholesky-ORHR (n=4096)	最大の行列に対してGram-Choleskyを使用
6	CUDA graph replay	カーネル起動のオーバーヘッドを削減
7	Fused V/T layout	スライス・コピーと一時変数を排除
8	Split16 panels	Gram-Schmidt法による末尾処理の最適化
9	Shape specialization	行と列のハードコード化および融合されたリダクション
10	Superpanels	V256/T256パックとdirect-H returns (1.80k $\mu$s)

重要な洞察とトレードオフ

ドメイン知識 vs. 自動化

エージェントは大幅な高速化を達成できる一方で、制御にはドメイン知識が不可欠でした。開発者は、上位10個のソリューションが以下のようにさらに最適化されていることに気づきました。

データの検出: 特定の入力分布（例：低ランクの場合）を活用。
ライブラリの除去: 残りのPyTorch関数（例：triangular solve）をカスタムCUDA/Triton実装に置きなる。
精度管理: 繰り返しの型変換を避けるため、行列の残りの部分をFP16に保持し続ける。

汎用性 vs. 特異性

コミュニティの議論では、この「ループ・エンジニアリング」手法における重大なリスクが浮え彫りになりました。**ベンチマークへの過学習（overfitting to the benchmark）**です。

"8 out of the 10 top solutions... completely broke at any other input than the competition ones. The only solutions that did not break... were made by experts who... followed and adjusted their solution in reasonable bounds."

これは、LLM駆動のループは、特定の、明確に定義されたベンチマークマークアップを解くために非常に強力ですが、人間による強力な制御なしには汎用的な堅牢性を維持することが困難であることを示唆しています。