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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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

SymbolDefinition
x∈Xx\in\mathcal{X}Kernel program implementation
o∈Oo\in\mathcal{O}Observation tuple from execution
E:X→O\mathcal{E}:\mathcal{X}\to\mathcal{O}Evaluator function
J:X→RJ:\mathcal{X}\to\mathbb{R}Scalar objective score
Ht\mathcal{H}_{t}History {(xi,oi)}i=1t\{(x_{i},o_{i})\}_{i=1}^{t}
StS_{t}Search state
Table 1: Notation and Definitions
Figure 1: Overview of K-Search . The framework operates on a Search State StS_{t} structured as a search tree. The tree consists of Closed nodes (blue, visited states with attached program like x12x_{12} ) and a Frontier of Open nodes (orange, pending hypotheses like u13u_{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 VV ; (2) Local Refinement , where a stochastic policy πcode\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 VV , e.g., u11u_{11} dropping from 0.9 to 0.6), and Prune (removing less promising nodes like u10u_{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.
KernelArch.Model ContextCharacteristics & Challenges
MLA Paged PrefillHopperDeepSeek-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 DecodeHopperDeepSeek-V3Targets large dynamic batches. Latency-bound at low batch sizes. Uses Persistent Data Layout and specialized Hopper-instructions to minimize memory movement.
GQA Paged DecodeHopperQwen3-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 MoEBlackwellDeepSeek-V3Challenges 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.
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 IDLang.ModelIter.Latency ( \SIUnitSymbolMicros\mathrm{\SIUnitSymbolMicro s} , ↓\downarrow )
shiyegaoCUDA––10741074
Zeyu ShenTriton––11401140
TTTTritonGPT-OSS-20B w/ RL25,60011611161
K-Search (Ours)TritonGPT-5.2 + Gemini-3-Pro3001030
Table 3: Top GPUMODE TriMul submissions on NVIDIA H100. Latency is the geometric mean across a fixed set of benchmark cases.
More in this topic

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

Explore similar work

CardsList
May 8, 2026

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

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.

Zixuan Huang, Da Chen, Kecheng Huang +5

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

GPU Kernel OptimizationGPU AccelerationEvolutionary OptimizationLarge Language Model-Guided OptimizationLLM-Based Program SynthesisGPU Kernel Generation
AskPDF
Mar 12, 2026

KernelFoundry: Hardware-aware evolutionary GPU kernel optimization

GPU kernel optimization challenges LLMs beyond standard coding tasks, as it requires an understanding of hardware architecture, parallel computing optimization strategies, and profiling outputs. However, most existing approaches leveraging LLMs for kernel generation apply standard prompting and feedback loops, considering hardware only through profiling feedback. We introduce KernelFoundry, an evolutionary framework that efficiently explores the space of GPU kernels through (1) MAP-Elites quality diversity search with kernel-specific behavioral dimensions to sustain exploration; (2) meta-prompt evolution that co-evolves prompts with kernels to uncover task-specific optimization strategies, and (3) a template-based parameter optimization approach to tune kernels to inputs and hardware. We evaluate this framework on Kernel-Bench, robust-kbench and custom tasks, generating SYCL kernels as a cross-platform GPU programming paradigm, and CUDA kernels for comparison to prior work. Our approach consistently outperforms the baseline methods and achieves an average speedup of 2.3 on KernelBench for SYCL. Moreover, KernelFoundry is implemented as a distributed framework with remote access to diverse hardware, allowing quick benchmarking and featuring a flexible user input layer to support kernel generation for a wide range of real use cases beyond benchmarking.

Nina Wiedemann, Quentin Leboutet, Michael Paulitsch +2

Intel Corporation

GPU Kernel OptimizationGPU AccelerationEvolutionary OptimizationLarge Language Model-Guided OptimizationLLM-Based Program SynthesisGPU Kernel Generation+1
AskPDF
Jun 25, 2026

EGG: An Expert-Guided Agent Framework for Kernel Generation

High-performance GPU kernels are critical for reducing the exponentially growing computational costs of large language models (LLMs), but their development heavily relies on manual tuning by domain experts. While recent advances in LLM-based approaches show promise for automating kernel generation, they still struggle to achieve both correctness and high performance. This limitation primarily arises from the lack of domain-specific optimization guidance, hindering effective exploration of the optimization space. We propose EGG, an Expert-Guided Agent Framework for Kernel Generation, which incorporates expert optimization principles to guide LLMs' decisions. Inspired by expert workflows, we decompose kernel generation into two hierarchical stages: 1) algorithmic structure design, which establishes a high-quality computational structure foundation; 2) hardware-specific tuning, which performs targeted adjustments through parallel mapping, tensor tiling, and memory optimization. This staged decomposition defines explicit optimization objectives, structuring the design space to achieve progressive refinement. To this end, a stage-aware multi-agent collaboration mechanism is designed for inter and intra-stage context management, ensuring stable optimization trajectories. Experiments on KernelBench and real-world workloads show that EGG achieves a 2.13x average speedup over PyTorch, outperforming existing agent-based and RL-based approaches.

Yaochen Han, Ke Fan, Hongxu Jiang +5

Beihang University, Beijing, China · Shanghai Jiao Tong University, Shanghai, China · Tsinghua University, Beijing, China +1

GPU Kernel OptimizationLarge Language Model-Guided OptimizationLLM-Based Program SynthesisGPU Kernel GenerationMulti-Agent LLMs
AskPDF