K-Search: Geração de Núcleo de LLM via Coevolução do Modelo de Mundo 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

Resumo

A otimização de kernels de GPU é crítica para sistemas modernos de aprendizagem de máquina eficientes, mas continua a ser um desafio devido à complexa interação de fatores de projeto e à rápida evolução do hardware. As abordagens automatizadas existentes normalmente tratam os Modelos de Linguagem de Grande Escala (LLMs) meramente como geradores estocásticos de código dentro de loops evolutivos guiados por heurísticas. Esses métodos frequentemente lutam com kernels complexos que exigem transformações estruturais multi-etapa coordenadas, pois carecem de capacidades de planeamento explícito e frequentemente descartam estratégias promissoras devido a implementações intermédias ineficientes ou incorretas. Para resolver isto, propomos a Busca via Modelo de Mundo Co-Evolutivo e construímos o K-Search com base neste método. Ao substituir heurísticas de busca estáticas por um modelo de mundo co-evolutivo, a nossa estrutura aproveita o conhecimento de domínio prévio dos LLMs para guiar a busca, explorando ativamente o espaço de otimização. Esta abordagem desacopla explicitamente o planeamento algorítmico de alto nível da instanciação de programa de baixo nível, permitindo que o sistema navegue por caminhos de otimização não monotónicos, mantendo-se resiliente a defeitos temporários de implementação. Avaliamos o K-Search em diversos kernels complexos do FlashInfer, incluindo kernels GQA, MLA e MoE. Os nossos resultados mostram que o K-Search supera significativamente os métodos de busca evolutiva state-of-the-art, alcançando uma melhoria média de 2.10x e um ganho de até 14.3x em kernels MoE complexos. Na tarefa GPUMode TriMul, o K-Search alcança desempenho state-of-the-art no H100, atingindo 1030us e superando tanto soluções evolutivas anteriores como soluções desenhadas 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.