K-Suche: LLM-Kernel-Generierung durch ko-evolutionäres intrinsisches Weltmodell
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
Die Optimierung von GPU-Kerneln ist entscheidend für effiziente moderne Machine-Learning-Systeme, bleibt jedoch aufgrund des komplexen Zusammenspiels von Designfaktoren und der rasanten Hardwareentwicklung eine Herausforderung. Bestehende automatisierte Ansätze behandeln Large Language Models (LLMs) typischerweise lediglich als stochastische Code-Generatoren innerhalb von heuristisch gesteuerten evolutionären Schleifen. Diese Methoden scheitern häufig an komplexen Kerneln, die koordinierte, mehrstufige strukturelle Transformationen erfordern, da ihnen explizite Planungsfähigkeiten fehlen und sie vielversprechende Strategien aufgrund ineffizienter oder falscher Zwischenimplementierungen verwerfen. Um dieses Problem zu lösen, schlagen wir Search via Co-Evolving World Model vor und implementieren darauf aufbauend K-Search. Indem wir statische Suchheuristiken durch ein ko-evolvierendes Weltmodell ersetzen, nutzt unser Framework das Vorwissen der LLMs über die Domäne, um die Suche zu steuern und den Optimierungsraum aktiv zu erkunden. Dieser Ansatz entkoppelt explizit die hochrangige algorithmische Planung von der niederrangigen Programminstanziierung, was es dem System ermöglicht, nicht-monotone Optimierungspfade zu navigieren und gleichzeitig resistent gegen temporäre Implementierungsdefekte zu bleiben. Wir evaluieren K-Search an verschiedenen komplexen Kerneln aus FlashInfer, darunter GKA-, MLA- und MoE-Kernel. Unsere Ergebnisse zeigen, dass K-Search state-of-the-art evolutionäre Suchmethoden signifikant übertrifft und durchschnittlich eine 2,10-fache Verbesserung sowie bis zu einen 14,3-fachen Gewinn bei komplexen MoE-Kerneln erzielt. Bei der GPUMode-TriMul-Aufgabe erreicht K-Search auf der H100 state-of-the-art-Leistung mit 1030us und übertrifft damit sowohl frühere evolutionäre als auch von Menschen entwickelte Lösungen.
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.