K-Search: LLM-kernelgeneratie via co-evoluerende intrinsieke wereldmodellen

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

February 22, 2026

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

cs.AI

Samenvatting

Het optimaliseren van GPU-kernels is cruciaal voor efficiënte moderne machine learning-systemen, maar blijft uitdagend vanwege het complexe samenspel van ontwerpfactoren en de snelle evolutie van hardware. Bestaande geautomatiseerde benaderingen behandelen Large Language Models (LLM's) doorgaans slechts als stochastische codegeneratoren binnen door heuristieken geleide evolutionaire lussen. Deze methoden worstelen vaak met complexe kernels die gecoördineerde, meerstaps structurele transformaties vereisen, omdat ze expliciete planningscapaciteiten missen en veelbelovende strategieën frequent verwerpen vanwege inefficiënte of incorrecte tussentijdse implementaties. Om dit aan te pakken, stellen wij Search via Co-Evolving World Model voor en bouwen we K-Search op basis van deze methode. Door statische zoekheuristieken te vervangen door een co-evoluerend wereldmodel, benut ons framework de voorafgaande domeinkennis van LLM's om de zoektocht te sturen en de optimalisatieruimte actief te verkennen. Deze aanpak ontkoppelt expliciet hoogwaardige algoritmische planning van laagwaardige programma-instantiatie, waardoor het systeem niet-monotone optimalisatiepaden kan bewandelen en tegelijkertijd veerkrachtig blijft tegen tijdelijke implementatiedefecten. We evalueren K-Search op diverse, complexe kernels uit FlashInfer, waaronder GKA-, MLA- en MoE-kernels. Onze resultaten tonen aan dat K-Search aanzienlijk beter presteert dan state-of-the-art evolutionaire zoekmethoden, met een gemiddelde verbetering van 2.10x en een winst tot 14.3x op complexe MoE-kernels. Op de GPUMode TriMul-taak bereikt K-Search state-of-the-art prestaties op de H100, met een tijd van 1030us, waarmee het zowel eerdere evolutionaire als door mensen ontworpen oplossingen overtreft.

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.