Codexによる自動リサーチ：カーネルを232倍高速化する方法

Codexによる自動リサーチ：カーネルを232倍高速化する方法 GPUモードのqr_v2問題でCodexをGPUモードで使用し、ベースラインから232倍高速化を達成 08 Jul, 2026 目次 はじめに コンテストの概要 問題の紹介 なぜこの問題は自動リサーチに適しているのか 学習：より良い質問をするために（任意） QR分解の数学：Householder反射 ブロック化されたHouseholderアルゴリズムの助けを借りて、逐次処理を小さくする その他の課題 Codexの最大活用 カーネルの進捗 画期的なアイデア 局所的極大値から脱するためのアイデア多様性の導入 実装のヒント もっとうまくできたこと 結論 参考文献 謝辞 はじめに コンテストの概要 GPU Modeは、Core Automationと協力して、最近、自動リサーチをテーマにしたコンテストを開催しました。問題は、バッチ処理された正方コンパクトHouseholder QR分解、すなわちQR分解を実装することでした。私は183人の参加者中12位に入賞し、ベースラインソリューションに対して232倍の高速化を達成しました。この記事では、私がどのようにしてそこに至ったかを説明します。アプローチ、学習、そして直面したボトルネックについて説明します。これは自動リサーチへの私の最初の真剣な試みでした。一部の人はこれを「ループエンジニアリング」と呼ぶでしょうが、正直なところ、それも構いません。このブログ記事のほとんどを理解するために、数学や問題自体を詳細に理解する必要はないことに注意してください。私は、数学や問題自体を二次的なものとして扱い、私のアプローチに焦点を当てました。なぜなら、これを読む人のほとんどはコンテストに参加していないからです。完全なコンテストページはこちらで確認できます：問題リンクとリーダーボード このコンテストは、GPU Modeの「線形代数カーネル：リサーチ時代」シリーズの一部でした。 問題の紹介 バッチ処理された正方FP32 CUDA行列A（形状：batch x n x n）が与えられ、torch.geqrf(A)と同じコンパクトHouseholder QR表現を返す必要がありました。これは、上三角部分がRで、下三角部分にHouseholderベクトルを格納したH行列と、リフレクター係数のタウベクトルです。チェッカーはtorch.linalg.householder_product(H, tau)でQを再構築し、R = triu(H)として、A≈QR、Q⊤Q≈I、Q⊤A≈Rを検証しました。正しい提出の中から、リーダーボードは形状と条件付けケース全体での実行時間の幾何平均によってランク付けされました。重要なサイズは、バッチ処理された512 x 512のような正方行列で、さらに大きな1024、2048、4096のケースもありました。低ビットのFP16、FP8、またはNVFP4の使用は内部的に許可されましたが、返された因子はFP32スタイルのQRチェックを満たす必要がありました。 3x3の小さな例は次のとおりです。 A=[12−5146167−68−424−41]=[6/7−69/175−58/1753/7158/1756/175−2/76/35−33/35]⏟Q[1421−140175−700035]⏟R ここで、Qは直交行列であり、その列は単位長で互いに垂直であり、Rは上三角行列であり、対角線より下のすべてがゼロであることを意味します。コンテストでは、密なQとRを直接出力するように求められていませんでした。代わりに、チェッカーがQを再構築し、上三角からRを読み取ることができるコンパクトなHouseholder形式が求められました。上記の3x3の例では、最初の反射が最初の列（12、6、-4）を一度に（-14、0、0）にまっすぐにマッピングします。-14がR11になります。それがどのように機能するかは、数学のセクションにあります。 なぜこの問題は自動リサーチに適しているのか GPU Modeは、エージェントフレンドリーなpopcorn CLIを参加者に提供しています。エージェントはこれを使用して、テスト、ベンチマーク、およびリーダーボードへの直接提出が可能です。チェッカーは、全体的な実行時間の幾何平均とともに、形状ごとのフィードバックも提供しました。鋭い観察者は、これがループを作成するのに適したセットアップであることに気づくでしょう。エージェントは、タイトなフィードバックループを熱望しています。それらは、エージェントが好きなだけヒルクライムすることを可能にします。GPU Modeのコンテストでは、通常、カーネルを反復処理する方法が提供されます。直接提出するか、Modalのようなスポンサーがクレジットを提供します。ここでは、主催者は間隔を空ける限り、実質的に無制限の提出を許可しました。そうしないと、キューが長くなり、すべての実行がタイムアウトしました。ある時点で、誰もが提出を連打していたため、ワークスペースはModalクレジットを使い果たしました。学習をアクセスしやすくするための良い方法です。14日間で、私は1500回以上の提出を行いました。 学習：より良い質問をするために 私はGPUカーネル最適化の基本（主にTritonで、CUDAの理解も少しあります）を1年間知っていましたが、この分野で専門的に働いたことはありませんでした。私が伝えたいのは、私はリーダーボード上の周りの人々の中でアンダードッグだったということです。リーダーボードで私のすぐ上にいた人物（CUDA Colonel）は、NVIDIAのプリンシパルエンジニアです。Anyway aura farming aside、基本的なことを知っており、最近GatedDeltaNetについて読んだばかりだったので、一般的なGPUカーネルの専門用語には慣れていました。未知の未知を既知の未知に変換するため、LLMをより良くプロンプトできるほど、何かをより良く知ることができます。同時に、このコンテストはドメイン知識なしでも実行可能であったことに注意する価値があります。おそらくトップ10に入ることはできないでしょうが、ハーネス/エージェントループなどに依存するだけで、ベースラインよりも大幅な高速化を得ることができます。コンテストでの私の最初のステップは、QR分解とは何か、そしてそれがどのように行われるかを学ぶことでした。Gram-SchmidtやHouseholder反射など、さまざまな方法があります。コンテストではHouseholder反射が義務付けられていました。私はClaudeと何度もやり取りし、いくつかのYouTube動画を見て直感を築きました。Claudeとの議論の後、ブロック化されたHouseholderアルゴリズムをメインアーキテクチャとして、後続のWY更新を使用する必要があることが明らかになりました。GPT-5.5もこれについて良いアイデアを持っていました。QR分解はかなりよく知られた問題です。行列分解は、特に行列前処理を使用する方法（Shampooスタイルのオプティマイザーや関連アプローチなど）において、いくつかの最新のオプティマイザーバリアントで現れるため、興味深い概念だと感じました。Muon（Kimiによって使用される）も良い例です。重み更新を1つの巨大なフラット化されたベクトルとして扱うのではなく、行列構造を維持し、通常は数回のNewton-Schulzイテレーションを介して極分解を近似することで、モーメンタム更新を直交化します。 （任意）QR分解の数学：Householder反射 数学に興味がある場合は、このセクションをざっと読むことをお勧めします。それ以外の場合は自由にスキップしてください。注意すべき唯一の点は、Householder QRには逐次的な依存関係があり、GEMMを実行するのが問題になるということです。ブロック化されたHouseholderを使用して、より行列乗算の形状にします。 契約 契約を簡単にレビューします。入力はバッチ処理された正方FP32行列Aです。出力は、torch.geqrfが返すコンパクトな（H、tau）形式です。Hの上三角部分はRです。対角線の下には、HがHouseholderベクトルを格納し、tauは列ごとに1つのスカラーを格納します。チェッカーは（H、tau）からQを再構築し、A ≈ QRを検証します。 鏡 一時的に行列を忘れてください。バスルームの鏡では、あなたの反射はガラスの向こう側で、ガラスの正面にいるのと同じくらい遠くにあります。x⟂がガラスに垂直に突き出ている部分である場合、反射は単にその部分を2回引きます：xreflected=x−2x⟂ したがって、Householder反射は、垂直な部分を見つけてそれを2回引くことです。 鏡の保存 Householderベクトルは鏡であり、コンパクトに保存されます。コードでは、鏡面全体を持ち歩きません。鏡からまっすぐ伸びる1つのベクトルvを保存します。鏡はvに垂直なすべてのものですが、反射はvに沿って移動します。垂直な部分は、vに沿ったxの影、つまりv⊤xv⊤v個のvのコピーです。これを上記の減算に代入します：ℋx=x−τv(v⊤x),τ=2v⊤v したがって、tauは単に2v⊤vです。つまり、2の係数とvの長さが1つの事前計算された数値にバンドルされています。vが鏡を選択し、tauが更新をスケーリングします。（数学的なリフレクターをℋjと書き、Hをコンパクトな出力行列として予約します。） 2DでのHouseholder反射 tau = 0.00 x reflected x v x_parallel x_perp 鏡のルール：x_parallelを維持し、x_perpを反転させます。したがって、反射されたx = x - 2x_perpです。簡略化された2D Householderステップ。オレンジ色のベクトルをドラッグします。紫色の鏡が変わり、緑色の反射ベクトルが水平軸に着地します。