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

Shiyi Cao; Ziming Mao; Joseph E. Gonzalez; Ion Stoica

cs.AIFeb 22, 2026

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

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

Organizations: UC Berkeley

Abstract

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.

Figures & tables

Symbol	Definition
$x\in\mathcal{X}$	Kernel program implementation
$o\in\mathcal{O}$	Observation tuple from execution
$\mathcal{E}:\mathcal{X}\to\mathcal{O}$	Evaluator function
$J:\mathcal{X}\to\mathbb{R}$	Scalar objective score
$\mathcal{H}_{t}$	History $\{(x_{i},o_{i})\}_{i=1}^{t}$
$S_{t}$	Search state

Table 1: Notation and Definitions

Figure 1: Overview of K-Search . The framework operates on a Search State $S_{t}$ structured as a search tree. The tree consists of Closed nodes (blue, visited states with attached program like $x_{12}$ ) and a Frontier of Open nodes (orange, pending hypotheses like $u_{13}$ ). The workflow iterates through three phases: (1) Action Selection , where the most promising action node is retrieved from the frontier based on world model estimated priority score $V$ ; (2) Local Refinement , where a stochastic policy $\pi_{\text{code}}$ samples concrete implementations until stagnation; and (3) World Model Update , where the LLM reasons over the trajectory to update the search tree via Insert (adding new actions), Update (adjusting $V$ , e.g., $u_{11}$ dropping from 0.9 to 0.6), and Prune (removing less promising nodes like $u_{10}$ ).

Figure 2: K-Search Search Trace Visualization. It tracks the evolution of the Search State across search rounds on the MLA Paged Decode kernel (refer to Section 4 for setup details). A round corresponds to one candidate program evaluation. Nodes represent actions (dark=Closed, light=Open), annotated with their instantiated program performance (closed nodes) or priority scores (open nodes). The timeline highlights how the kernel is improved and how the LLM dynamically Inserts new hypotheses, Updates beliefs, and Prunes less promising branches based on evolved understanding.

Kernel	Arch.	Model Context	Characteristics & Challenges
MLA Paged Prefill	Hopper	DeepSeek-V3 Liu et al. (2024)	Implements attention with paged KV cache and split ckv (compressed, no RoPE) and kpe (RoPE). Bandwidth-bound in long contexts; requires matrix absorption and Split-K optimization.
MLA Paged Decode	Hopper	DeepSeek-V3	Targets large dynamic batches. Latency-bound at low batch sizes. Uses Persistent Data Layout and specialized Hopper-instructions to minimize memory movement.
GQA Paged Decode	Hopper	Qwen3-A3B-30B Yang et al. (2025)	Memory bound. Optimizations include fusing multiple query heads per KV head to reuse KV loads, and using vectorized/async paged-KV gather with pipelined global to shared memory prefetch.
FP8 MoE	Blackwell	DeepSeek-V3	Challenges include irregular data-dependent routing, load balancing, and managing FP8 packing/scaling overhead.

Table 2: Representative kernels used for evaluation. Detailed configurations (e.g., head counts) are omitted for brevity. We provide example optimization challenges based on FlashInfer flashinfer-ai (2026) implementations.

Panel of Figure 3 — Figure 3: Main Results (3 runs each). (a) compares the kernels best-so-far scores generated by the three systems across 120 iterations. (b) provides a per-workload analysis for all compared systems. (c) shows the fraction of workloads for which the best kernel from each system achieves the specified speedup over the FlashInfer baseline.

Submission ID	Lang.	Model	Iter.	Latency ( $\mathrm{\SIUnitSymbolMicro s}$ , $\downarrow$ )
shiyegao	CUDA	–	–	$1074$
Zeyu Shen	Triton	–	–	$1140$
TTT	Triton	GPT-OSS-20B w/ RL	25,600	$1161$
K-Search (Ours)	Triton	GPT-5.2 + Gemini-3-Pro	300	1030

Table 3: Top GPUMODE TriMul submissions on NVIDIA H100. Latency is the geometric mean across a fixed set of benchmark cases.

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 209 papers

Kernel Foundry: A Diagnosis-driven Evolutionary Kernel Optimizer with Multi-Experts
May 8, 2026Zixuan Huang, Da Chen, Kecheng Huang +5GPU Kernel Optimization GPU Acceleration
1The Chinese University of Hong Kong · 2Noah’s Ark Lab, Huawei
KernelFoundry: Hardware-aware evolutionary GPU kernel optimization
Mar 12, 2026Nina Wiedemann, Quentin Leboutet, Michael Paulitsch +2GPU Kernel Optimization GPU Acceleration
Intel Corporation
EGG: An Expert-Guided Agent Framework for Kernel Generation
Jun 25, 2026Yaochen Han, Ke Fan, Hongxu Jiang +5GPU Kernel Optimization Large Language Model-Guided Optimization
Beihang University, Beijing, China · Shanghai Jiao Tong University, Shanghai, China · Tsinghua University, Beijing, China +1

May 8, 2026J/K move · Enter open · S save

Kernel Foundry: A Diagnosis-driven Evolutionary Kernel Optimizer with Multi-Experts

Zixuan Huang, Da Chen, Kecheng Huang, Lihao Yin+4

1The Chinese University of Hong Kong · 2Noah’s Ark Lab, Huawei

GPU Kernel Optimization GPU Acceleration Evolutionary Optimization Large Language Model-Guided Optimization LLM-Based Program Synthesis GPU Kernel Generation

Open paper AskPDF arXiv

Generating high-performance GPU kernels remains challenging due to the need for both correctness and hardware-aware optimization. While large language models (LLMs) show promise in code generation, they often fail to produce kernels that are both correct and efficient. We propose Kernel Foundry, a diagnosis-driven evolutionary framework for automatic GPU kernel optimization. Our method combines expert-guided, retrieval-augmented initialization with a multi-island evolutionary search, where candidate kernels are iteratively refined using structured diagnostic feedback. A centralized experience library accumulates reusable optimization knowledge to guide subsequent evolution, while explicit mechanisms prevent cheating behaviors that bypass kernel-level computation. Experiments on KernelBench show that our method consistently improves both correctness and performance over strong baselines, achieving up to 100% correctness on Level~2.

Abstract

GPU Kernel Optimization

Explore similar work

Kernel Foundry: A Diagnosis-driven Evolutionary Kernel Optimizer with Multi-Experts

KernelFoundry: Hardware-aware evolutionary GPU kernel optimization

EGG: An Expert-Guided Agent Framework for Kernel Generation