Codexで232倍高速なGPUカーネルを実現
原題: Auto-research with codex: How I achieved a 232x Faster Kernel
なぜ重要か
AIエージェントによる高性能GPUカーネルの自動探索が現実的な成果を上げており、HPC領域でのLLM活用の実用性を示す事例として注目される。
GPU ModeとCore Automationが共催した自動研究コンテストで、エンジニアのSankalp氏がOpenAIのCodexを活用したエージェント主導の最適化手法により、QR分解(qr_v2)カーネルのベースライン比232倍の高速化を達成し、183名中12位に入賞した。コンテストは2026年7月に実施された。
GPU ModeとCore Automationは、「Linear Algebra Kernels in the Age of Research」シリーズの一環として、自動研究をテーマにしたコンテストを開催した。課題はバッチ処理された正方行列(FP32 CUDA)に対するコンパクトHouseholder QR分解の実装で、torch.geqrfと同等の出力(上三角がRのH行列とtauベクトル)を返す必要があった。対象サイズは512×512を中心に、1024、2048、4096の正方行列が含まれ、正確性はQ⊤Q≈IおよびQR≈Aで検証された。内部演算にはFP16やFP8の利用も許可されたが、最終出力はFP32精度の基準を満たす必要があった。
Sankalp氏が採用したのは「auto-research(自動研究)」と呼ばれるアプローチで、Codexエージェントにテスト・ベンチマーク・リーダーボードへの提出までを自動化させた。GPU ModeはエージェントフレンドリーなCLIツール「popcorn」を参加者に提供しており、エージェントがこれを通じて反復的に改善を重ねる構造を取っていた。
主要なブレークスルーはブロック化Householderアルゴリズムによるシリアル処理の削減と、ローカル最適解から脱出するためのアイデア多様性の導入にあった。Codexに対しては、数学的背景を学習した上でより精度の高い質問を行う工夫も施された。最終的に幾何平均ランタイムで183名中12位、ベースライン比232倍の性能向上を達成した。同氏はこのアプローチを「ループエンジニアリング」とも呼ばれる手法だと説明しつつ、改善余地として「より早い段階でのアイデア多様化」を挙げている。