Codexで自動リサーチ！カーネルの処理速度を232倍に爆速化した話

Codexを活用してカーネルを自動最適化した結果、処理速度を232倍まで高速化することに成功しました。このプロジェクトの全貌と、AIによる自動探索アプローチのノウハウを共有します。

ここ数日、新しい決定版となるDeepSeek v4を試してみたんだ。開発が半ば放置されている動画圧縮コーデックのリポジトリを与えて、「ベンチマーク→プロファイル→検証→リサーチ→改善」といういつものループを実行させてみたよ。わざわざこのコーデックを選んだのは、作者がビットストリームの検証ツールを用意していて、独自実装を試す際に壊していないか確認できるからだ。エージェントにはコンパイラのプロファイラと、素晴らしい出力を出してくれるIntelのVTuneへのアクセス権を与えた。数時間でLLMは圧縮・伸張アルゴリズムのSSEおよびAVX実装を生成し、シングルコアでのパフォーマンスをほぼ倍増させてしまった。さらに、NVIDIAのNSIGHTプロファイラをガイドとしてCUDA実装も作らせたんだけど、これもかなりいい感じに動き始めている。

個人的には、LLMはPrologや線形計画法の高度なバージョンとして扱うべきだと思うんだ。制約を与え、正しさを検証する手段を用意し、明確なゴールを示す。もしLLMが自己検証と軌道修正を行えるなら、基本的には自動操縦で放置できるはずだよ。

マジかよ！ソロのエンジニアがこんなことを実現できるなら、OAIやAnthropicの周りで騒がれていることなんて、かなり大したことないように見えてくるな。

GPUカーネルとSIMDに関する学習データが特に充実しているみたいだな。

これってモデルの研究者たち自身に役立つから力を入れているのか、それとも単に言語モデルが非常に得意かつ人間には難しいサブドメインだからなのか、どっちなんだろうね？

投稿の中には実際のルールを無視しているものがいくつかあるな。edinhの1位のやつなんて「bypass ban check（BAN回避）」なんていう一文が入っているし。

コレスキー分解って、ある地点でハウスホルダー変換の代わりに使われるけど、場合によってはより高速でも安定性に欠けるんじゃなかったっけ？今年試験のためにハウスホルダーを使ったQR分解の小プロジェクトをやった時の記憶を辿っているだけなんだけどね。つまり、もし標準的なtorchの演算より速いなら、それがデフォルトの標準になっていない何かしらの正当な理由があるはずだよ。間違っているかもしれないけど、確信は持てないな。

メタな感想だけど、AIが生成したようには見えない長文を読むのは新鮮でよかった。ありがとう。

CPU+GPU向けの初のOSS組み込み可能なCypherプロパティグラフクエリエンジンであるGFQLのカスタムバリアントを作るのは非常に興味深い体験だったよ。

• polarsのような新しいバックエンドのローンチを加速させたし、新しい遅延モードやプランナーなど、根本的に新しいパスも取り入れた

• 当初はGPUのベンチマークスコアでトップを狙っていたけど、今ではCPUスコアでもトップを維持できている！

長期的には、クエリエンジンであることの意味を再考させてくれる点が一番面白いと思っている。現時点では、我々のユースケースや主要な業界ベンチマーク、そしてユーザーからのワークロードで、全般的に最速のものを作っている最中だ。同時に、JITやマルチステージコンピューティングと同じように、カスタムインデックスを埋め込むよりも面白い、ユーザーができる新しい事前最適化技術を模索している。本質的に、我々のエージェントが高速な最適化を実行できるなら、ユーザー側のエージェントにも公開できる安全なフックを用意するべきだろう！

このコンテストで注目すべきなのは、トップ10のうち8つものソリューションが、この手法で最適化されていたにもかかわらず、コンテスト用以外の入力では完全に動かなくなったということだ。

OOD（Out-of-Distribution）形状でテストしても壊れなかった唯一のソリューションは、GPUプログラミングを深く理解している専門家によるものだった。彼らは2万5千行ものCUDAコードを生成するのではなく、適切な範囲内で自身のソリューションを調整していたんだ。

ここから得られる教訓は、この種のアプローチは常に特定用途のための最適解を出すものだということであって、モデルを汎用的な解決策へ導くのは遥かに難しいタスクだということだよ。だから、特定のモデル形状に対する推論プロバイダーをしているなら最高だ、どんどんやるべきだ。でも、オープンソースライブラリのメンテナーなら、これはあまり役に立たないね。

私の経験と重なるな。LLMは最適化が本当に得意で、ほとんどの人間よりも優れている。ただ、人が何らかの努力で最適化を詰め込んだ究極のピークパフォーマンスには届かない傾向がある。

とはいえ、ほとんどの問題はほとんど最適化されていないのが現状だから、それでも大抵の場合は大きな成果だと言えるね。

パフォーマンス最適化ループには、次のプロセスがかなりうまくいくことがわかったよ。

• 100%のパス網羅（単なるステートメント網羅ではなく）に達するまでエージェントにユニットテストを生成させる。すべての数値テストは退行バグを防ぐためにゴールデン値との比較チェックを行うこと。

• 手元にある最も広範で代表的なエンドツーエンドのテストケースを使って、パフォーマンス改善ループを全力で回させる。その際、フレイムグラフを生成させて、必要に応じて介入や軌道修正ができるようにしておく。

• 出力値に対して必要なら1 ULP程度の誤差を許容するようにする。ビット単位で完全に一致させることに固執して自滅するのを防ぐためだ。

• コードにはバグが残っている可能性が高いので、進行に合わせて正当性のエラーをフラグ立てさせる。

これは私がPythonからRustへ言語移植をする際にも使っている手法で、鉄壁のテストカバレッジがあれば、コードが元の意図から乖離するのを防いでくれるよ。