arXiv cs.AIは2026年6月2日(現地時間)、「Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels」と題する研究を発表した。この研究は、大規模言語モデル (LLM) を利用してCUDAカーネルの生成と最適化を行うエージェントハーネスシステム「Kernel Forge」を提示している。これにより、機械学習モデルのランタイムにおけるGPUコード最適化の労力軽減を目指す。

機械学習モデルのランタイムの大部分は、行列乗算、畳み込み、正規化といった少数の計算カーネルに費やされている。これらのカーネルを最適化することは、レイテンシーとコストを削減する最も直接的な方法の一つであるものの、従来は専門家が低レベルのGPUコードを手書きする必要があった。

既存のLLMベースの最適化ツールについては、ランダムに生成されたテンソルと独立したカーネルで評価されることが多く、開発者が手動で再統合する必要があるスタンドアロンのCUDAコードを出力すること、主にLLM PyTorchモデルのみを対象としていること、および結果の検査やデバッグのサポートが限定的であることが指摘されている。

「Kernel Forge」は、変更されていない任意のPyTorchモデルをそのまま受け入れるオープンソースのエンドツーエンドエージェントハーネスとして開発された。このシステムは、ビジョン、拡散、およびLLMワークロードをサポートし、単一の線形な改良チェーンではなく、複数の最適化経路を探索するためにモンテカルロ木探索 (MCTS) を採用する。また、進行状況の監視、候補カーネルの検査、障害のデバッグのためのグラフィカルユーザーインターフェースを備えている。

研究チームは、NVIDIA DGX SparkとGB10 GPUを使用し、ビジョン、拡散、LLMワークロードをカバーする4つのPyTorchモデルで「Kernel Forge」を評価した。各カーネルあたりわずか50回の最適化イテレーションで、14個のカーネルがPyTorchの eager mode を上回る性能を示した。具体的には、ResNet-50の adaptive_avgpool2d で1.52倍、Stable Diffusion 3.5 Mediumの group_norm で1.70倍、Gemma 4 E2Bの softmax で2.83倍、Qwen 3.5 35B-A3Bの softmax で1.54倍の性能向上が達成された。


参考: arXiv cs.AI (アーカイブ) — 2026年7月29日 13:00 (JST)

原文ハイライト

"Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels"

この記事をシェア
X はてブ LinkedIn