K-Search: Generación de Núcleos de LLM mediante la Co-evolución de un Modelo Mundial Intrínseco

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

February 22, 2026

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

cs.AI

Resumen

La optimización de kernels de GPU es fundamental para sistemas modernos de aprendizaje automático eficientes, pero sigue siendo un desafío debido a la compleja interacción de factores de diseño y la rápida evolución del hardware. Los enfoques automatizados existentes suelen tratar a los Modelos de Lenguaje a Gran Escala (LLMs) meramente como generadores estocásticos de código dentro de bucles evolutivos guiados por heurísticas. Estos métodos a menudo tienen dificultades con kernels complejos que requieren transformaciones estructurales coordinadas y multi-paso, ya que carecen de capacidades de planificación explícita y frecuentemente descartan estrategias prometedoras debido a implementaciones intermedias ineficientes o incorrectas. Para abordar esto, proponemos la Búsqueda mediante un Modelo Mundial en Co-evolución y construimos K-Search basado en este método. Al reemplazar las heurísticas de búsqueda estáticas con un modelo mundial en co-evolución, nuestro marco aprovecha el conocimiento de dominio previo de los LLMs para guiar la búsqueda, explorando activamente el espacio de optimización. Este enfoque desacopla explícitamente la planificación algorítmica de alto nivel de la instanciación de programas de bajo nivel, permitiendo que el sistema navegue por trayectorias de optimización no monótonas mientras se mantiene resiliente a defectos temporales de implementación. Evaluamos K-Search en diversos kernels complejos de FlashInfer, incluyendo kernels GQA, MLA y MoE. Nuestros resultados muestran que K-Search supera significativamente a los métodos de búsqueda evolutiva más avanzados, logrando una mejora promedio de 2.10x y una ganancia de hasta 14.3x en kernels MoE complejos. En la tarea GPUMode TriMul, K-Search alcanza un rendimiento de vanguardia en H100, llegando a 1030us y superando tanto a soluciones evolutivas previas como a las diseñadas por humanos.

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.