K-サーチ：内在的世界モデルの共進化によるLLMカーネル生成

要旨

GPUカーネルの最適化は、現代の効率的な機械学習システムにおいて極めて重要であるが、設計要因の複雑な相互作用とハードウェアの急速な進化により、依然として困難な課題である。既存の自動化アプローチは、典型的には大規模言語モデル（LLM）を、ヒューリスティックに導かれた進化的ループ内での確率的コード生成器として扱うことが多い。これらの手法は、協調的な多段階の構造変換を必要とする複雑なカーネルに対して苦戦することが多い。なぜなら、明示的な計画能力を欠き、非効率的または不正確な中間実装のために有望な戦略を頻繁に破棄してしまうからである。この問題に対処するため、我々は「共進化する世界モデルによる探索（Search via Co-Evolving World Model）」を提案し、この手法に基づいてK-Searchを構築した。静的な探索ヒューリスティックを共進化する世界モデルで置き換えることで、本フレームワークはLLMの事前の領域知識を活用して探索を導き、最適化空間を能動的に探査する。このアプローチは、高レベルのアルゴリズム的計画と低レベルのプログラム具体化を明示的に分離し、システムが非単調な最適化経路を進みながら、一時的な実装上の欠陥に対して耐性を保つことを可能にする。我々はK-Searchを、FlashInferのGQA、MLA、MoEカーネルを含む多様で複雑なカーネルで評価した。その結果、K-Searchは最新の進化的探索手法を大幅に上回り、平均2.10倍、複雑なMoEカーネルでは最大14.3倍の性能向上を達成した。GPUMode TriMulタスクでは、K-SearchはH100上で1030usを達成し、従来の進化的手法および人間が設計した解法の両方を凌駕する、最新の性能を実現した。

English

Optimizing GPU kernels is critical for efficient modern machine learning systems yet remains challenging due to the complex interplay of design factors and rapid hardware evolution. Existing automated approaches typically treat Large Language Models (LLMs) merely as stochastic code generators within heuristic-guided evolutionary loops. These methods often struggle with complex kernels requiring coordinated, multi-step structural transformations, as they lack explicit planning capabilities and frequently discard promising strategies due to inefficient or incorrect intermediate implementations. To address this, we propose Search via Co-Evolving World Model and build K-Search based on this method. By replacing static search heuristics with a co-evolving world model, our framework leverages LLMs' prior domain knowledge to guide the search, actively exploring the optimization space. This approach explicitly decouples high-level algorithmic planning from low-level program instantiation, enabling the system to navigate non-monotonic optimization paths while remaining resilient to temporary implementation defects. We evaluate K-Search on diverse, complex kernels from FlashInfer, including GQA, MLA, and MoE kernels. Our results show that K-Search significantly outperforms state-of-the-art evolutionary search methods, achieving an average 2.10x improvement and up to a 14.3x gain on complex MoE kernels. On the GPUMode TriMul task, K-Search achieves state-of-the-art performance on H100, reaching 1030us and surpassing both prior evolution and human-designed solutions.

K-サーチ：内在的世界モデルの共進化によるLLMカーネル生成

K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model

要旨

Support