K-Search: Generazione del Kernel LLM tramite Co-Evoluzione del Modello Mondiale Intrinseco

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

February 22, 2026

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

cs.AI

Abstract

L'ottimizzazione dei kernel GPU è fondamentale per sistemi di machine learning moderni ed efficienti, ma rimane impegnativa a causa del complesso intreccio di fattori progettuali e della rapida evoluzione dell'hardware. Gli approcci automatizzati esistenti tipicamente trattano i Large Language Model (LLM) semplicemente come generatori stocastici di codice all'interno di loop evolutivi guidati da euristiche. Questi metodi spesso incontrano difficoltà con kernel complessi che richiedono trasformazioni strutturali coordinate e multi-step, poiché mancano di capacità di pianificazione esplicita e scartano frequentemente strategie promettenti a causa di implementazioni intermedie inefficienti o errate. Per affrontare questo problema, proponiamo la Ricerca tramite Modello del Mondo Co-Evolvente e sviluppiamo K-Search basandoci su questo metodo. Sostituendo le euristiche di ricerca statiche con un modello del mondo co-evolvente, il nostro framework sfrutta la conoscenza pregressa del dominio degli LLM per guidare la ricerca, esplorando attivamente lo spazio di ottimizzazione. Questo approccio dissocia esplicitamente la pianificazione algoritmica di alto livello dall'istanziazione di programma di basso livello, consentendo al sistema di navigare percorsi di ottimizzazione non monotoni rimanendo resiliente ai difetti temporanei di implementazione. Valutiamo K-Search su kernel complessi e diversificati provenienti da FlashInfer, inclusi i kernel GQA, MLA e MoE. I nostri risultati mostrano che K-Search supera significativamente i metodi di ricerca evolutiva all'avanguardia, ottenendo un miglioramento medio di 2.10x e un guadagno fino a 14.3x sui kernel MoE complessi. Sul task GPUMode TriMul, K-Search raggiunge prestazioni all'avanguardia su H100, toccando 1030us e superando sia le soluzioni evolutive precedenti che quelle progettate da umani.

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.