K-Search : Génération de noyaux LLM par co-évolution d'un modèle du monde intrinsèque

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

February 22, 2026

papers.authors: Shiyi Cao, Ziming Mao, Joseph E. Gonzalez, Ion Stoica

cs.AI

papers.abstract

L'optimisation des noyaux GPU est cruciale pour l'efficacité des systèmes modernes d'apprentissage automatique, mais reste difficile en raison de l'interaction complexe des facteurs de conception et de l'évolution rapide du matériel. Les approches automatisées existantes traitent généralement les grands modèles de langage (LLM) comme de simples générateurs de code stochastiques dans des boucles évolutives guidées par des heuristiques. Ces méthodes peinent souvent avec les noyaux complexes nécessitant des transformations structurelles multi-étapes coordonnées, car elles manquent de capacités de planification explicite et rejettent fréquemment des stratégies prometteuses en raison d'implémentations intermédiaires inefficaces ou incorrectes. Pour résoudre ce problème, nous proposons la recherche par co-évolution d'un modèle du monde et développons K-Search sur cette base. En remplaçant les heuristiques de recherche statiques par un modèle du monde co-évolutif, notre cadre exploite les connaissances préalables du domaine des LLM pour guider la recherche, en explorant activement l'espace d'optimisation. Cette approche découple explicitement la planification algorithmique de haut niveau de l'instanciation de programme de bas niveau, permettant au système de naviguer sur des chemins d'optimisation non monotones tout en restant résilient aux défauts d'implémentation temporaires. Nous évaluons K-Search sur divers noyaux complexes de FlashInfer, incluant les noyaux GQA, MLA et MoE. Nos résultats montrent que K-Search surpasse significativement les méthodes de recherche évolutive de pointe, obtenant une amélioration moyenne de 2,10x et jusqu'à un gain de 14,3x sur les noyaux MoE complexes. Sur la tâche GPUMode TriMul, K-Search atteint des performances de pointe sur H100, parvenant à 1030us et surpassant à la fois les solutions évolutives antérieures et celles conçues par l'homme.

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.