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

2026-02-22

AI Summary

Optimizing GPU kernels remains highly challenging due to design complexity and rapid hardware evolution, and existing LLM-based approaches lack explicit planning capabilities, struggling with complex scenarios requiring multi-step, coordinated structural transformations. This work proposes K-Search, a novel framework that introduces, for the first time, a co-evolving internal world model to decouple high-level algorithmic planning from low-level program implementation. By leveraging LLMs’ domain priors to guide exploration along non-monotonic optimization paths and tolerating transient implementation flaws, K-Search achieves significant performance gains. It delivers an average 2.10× speedup across FlashInfer kernels—including GQA, MLA, and MoE—with MoE attaining up to 14.3× acceleration. Furthermore, it reaches 1030 μs on the GPUMode TriMul task, surpassing both state-of-the-art automated and hand-optimized solutions.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• GPU kernel optimization
• Large Language Models
• evolutionary search
• code generation
• hardware evolution

Innovation

• K-Search
• co-evolving world model
• LLM-guided optimization
• GPU kernel generation
• non-monotonic search