Source-linked AI summary
K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model
Shiyi Cao, Ziming Mao, Joseph E. Gonzalez, Ion Stoica
TL;DR
Optimizing GPU kernels is difficult because large design spaces, rapid hardware evolution, and costly testing complicate manual and automated search, while existing evolutionary methods lack explicit planning for multi-step transformations. K-Search introduces a co-evolving LLM world model that separates optimization intent from implementation and uses execution feedback to guide search. It achieves a 2.10× average improvement over OpenEvolve, up to 14.3× on MoE, and state-of-the-art GPUMode TriMul performance at 1030 µs on H100.
Problem
GPU-kernel optimization involves complex design choices and changing hardware, while existing evolutionary approaches lack explicit planning for multi-step transformations and may discard strategies after temporary implementation defects.
Method
K-Search uses a co-evolving LLM world model to prioritize high-level optimization intents, assimilate execution feedback, and decouple planning from low-level program instantiation.
Results
K-Search achieves a 2.10× average improvement over OpenEvolve across diverse FlashInfer kernels, up to 14.3× on MoE, and 1030 µs on H100 for GPUMode TriMul.
Takeaways & Limitations
The results support using LLMs as planning-oriented intrinsic world models for complex GPU-kernel optimization rather than only as code generators.
Takeaways & Limitations
The evaluated evolutionary methods can prematurely discard theoretically sound strategies because intermediate edits may fail to improve the objective or compile temporarily.
Abstract
from arXiv · showhide
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.
1 Introduction
K-Search addresses the difficulty of optimizing complex GPU kernels by using a co-evolving LLM world model to plan high-level transformations separately from program instantiation. Across FlashInfer workloads, it substantially outperforms evolutionary baselines and reaches state-of-the-art performance on GPUMode TriMul.
- 1 Introduction: GPU kernel optimization is difficult because it combines large design spaces, rapidly changing hardware, and costly implementation, compilation, and profiling cycles.Relevant choices include tiling, memory layout, synchronization, and architecture-specific primitives.
- 1 Introduction: Existing LLM evolutionary methods treat language models primarily as stochastic code generators guided by heuristic search in program space.They select and mutate candidate programs directly, despite kernels often requiring coordinated structural transformations.
- 1 Introduction: Multi-step kernel optimizations can require intermediate edits that do not immediately improve performance, while temporary compilation errors may cause promising strategies to be discarded.This limits evolutionary methods’ ability to discover deep structural optimizations.
- 1 Introduction: K-Search uses a co-evolving LLM world model to maintain the search frontier, prioritize high-level optimization intents, and update its strategy from execution feedback.The framework decouples high-level planning from low-level program instantiation, supporting non-monotonic paths and temporary implementation defects.
- 1 Introduction: 2.10× average improvement over OpenEvolve and 14.3× improvement on the MoE kernel demonstrate K-Search’s gains across complex FlashInfer workloads.The evaluation covers GQA, MLA, and MoE kernels; K-Search also reaches state-of-the-art GPUMode TriMul performance at 1030 µs on H100.
2 Related Work
Prior work spans specialized GPU-kernel libraries, compiler and DSL-based optimization, iterative LLM refinement, evolutionary program search, and LLMs used as world models for planning. These lines of research establish the components that K-Search combines for kernel optimization.
- 2 Related Work: Specialized libraries target workload-specific GPU optimization, including FlashAttention for dense attention, FlashLinearAttention for attention variants, and FlashInfer for high-throughput LLM serving.These libraries reflect substantial engineering effort toward optimized kernels for relatively narrow workload classes.
- 2 Related Work: Compiler autotuners such as TVM and Ansor search tensor-program scheduling spaces with learned cost models, while Triton and CuTe provide higher-level, architecture-aware kernel abstractions.These approaches automate search or expose abstractions for explicit control over tiling and layouts.
- 2 Related Work: LLM-based GPU-kernel systems commonly use iterative generation or refinement driven by compilation, execution, and profiling feedback, with some adding evolutionary exploration.EvoEngineer is cited as an example of augmenting this paradigm with evolutionary strategies.
- 2 Related Work: Evolutionary program-search systems pair LLM-generated edits with execution-based evaluation and population or database management to improve programs.FunSearch and AlphaEvolve illustrate this pattern in mathematical, combinatorial, and codebase-evolution settings.
- 2 Related Work: World-model research uses LLMs for planning and decision making, including approaches that frame reasoning as planning or induce structured domain models.Recent work also examines LLMs as world models for planning beyond direct code generation.
3 K-Search
K-Search formulates GPU-kernel synthesis as budgeted search over a structured tree, using an LLM world model to plan optimization intents, instantiate implementations, and update the search state from execution outcomes. Its co-evolving design separates high-level planning from local code refinement, allowing promising strategies to survive temporary implementation defects and enabling adaptive insertion, reprioritization, and pruning.
- Problem Setup: GPU-kernel synthesis is formulated as maximizing speedup J(x) over a fixed evaluation budget B.The evaluator returns performance and metadata, while the objective is defined relative to a reference SoTA baseline.
- Motivation: Existing program-space evolutionary search couples optimization intent with implementation, so transient syntax errors can discard theoretically sound strategies.These methods also struggle with multi-step transformations whose intermediate edits may not immediately improve performance.
- Search via Co-Evolving World Model: K-Search represents the search process as a tree whose world model maintains explored actions, a frontier of pending intents, and priority scores V.An action pairs a parent program with a specific optimization intent, such as resolving bank conflicts through padding.
- Search via Co-Evolving World Model: Each iteration selects the highest-priority frontier action, samples concrete implementations with πcode until stagnation, and evaluates the resulting programs.The local refinement loop resets its stagnation counter after improvement and increments it after failures.
- Search via Co-Evolving World Model: After refinement, the world model updates the search tree by inserting new actions, updating frontier priorities, and pruning infeasible or redundant branches.The state transition uses accumulated execution outcomes to refine the model’s beliefs and sharpen subsequent search.
- Case Study: MLA Paged Decode: In the MLA Paged Decode trace, a composable refinement applying sm_scale during Q loading eventually reaches the global optimum at round 102.The trace illustrates that split-K can be ineffective alone but effective when composed with a strong fusion kernel.
- Takeaway: K-Search avoids enumerating a massive sparse program space by starting from high-level intents, filtering coding noise through local refinement, and reallocating search toward promising directions.The world model co-evolves with optimization progress, allowing dead ends to be pruned and strategies to be repositioned.
4 Experiments
K-Search is evaluated against OpenEvolve and ShinkaEvolve on FlashInfer kernels using standardized correctness and benchmarking procedures. It achieves higher overall and per-workload performance, with gains attributed to planning, persistent search state, and targeted kernel designs, while remaining sensitive to workload characteristics.
- Setup: The evaluation uses fixed iteration budgets, repeated runs, standardized evaluators, correctness checks, and common CUDA benchmarking infrastructure.Each method runs for 120 iterations and is repeated three times; candidates must pass functional correctness tests before receiving a non-zero score.
- Setup: K-Search is compared with OpenEvolve and ShinkaEvolve on representative FlashInfer kernels using common workloads and baseline configurations.Experiments use three automated methods, identical input workloads, and default baseline configurations.
- Overall Performance: 56.13 overall average final score gives K-Search a 2.10× improvement over OpenEvolve and a 2.21× improvement over ShinkaEvolve across four kernels.The four kernels are GQA decode, MLA decode, MLA prefill, and MoE; K-Search also leads each reported kernel-level comparison.
- Best Kernel per-Workload Performance: Across 4 kernel types and 152 workload traces, K-Search achieves higher performance than the baselines on the vast majority of workloads.Some GQA decode workloads with small batch sizes are exceptions.
- Best Kernel Fastp Analysis: At speedup ≥0.50 on GQA decode, K-Search succeeds on 87.5% of workloads versus 50.0% for OpenEvolve and 39.6% for ShinkaEvolve.For MLA prefill at speedup ≥0.40, K-Search reaches 57.9% while neither baseline reaches that threshold.
- Key observations: K-Search’s persistent search state and optimization-intent tracking support more targeted hypotheses than direct program-space evolution.The analysis reports that ShinkaEvolve often produces invalid or failed programs, while OpenEvolve shows high per-iteration variance and struggles on MoE.
5 Conclusion
K-Search treats LLMs as intrinsic world models with latent planning capabilities, replacing static search heuristics with a co-evolving model that reasons over optimization strategies. Its results support this approach across complex kernels and the GPUMODE TriMul task.
- K-Search uses a co-evolving world model to replace static search heuristics and guide optimization-space exploration.The LLM maintains planning over the search space rather than serving only as a code generator.
- K-Search demonstrates that LLMs can function as intrinsic world models with latent planning capabilities.
- 2.1× average improvement across diverse complicated kernels and up to 14.3× gains on MoE validate the proposed paradigm shift.
- K-Search achieves state-of-the-art performance on the GPUMODE TriMul competition.
A.1 Other generated kernels analysis
K-Search improves generated kernels through workload-adaptive execution strategies, including GPU-side handling of variable-length batches, register-resident queries, deeper prefetching, and adaptive sequence splitting. These choices keep computation parallel and avoid unnecessary reductions or staging, although the evaluated specifications and comparisons remain workload-specific.
- Kernel specifications: All three systems generate CUDA kernels for the same MLA Paged Prefill and MLA Paged Decode specifications.The specifications include fixed attention dimensions and layouts for the respective kernels.
- MLA Paged Prefill: K-Search resolves variable-length batch boundaries on the GPU while assigning each thread block a contiguous tile of 16 rows.Tiles that straddle sequence boundaries are handled on the fly using the sequence prefix-sum array.
- MLA Paged Prefill: K-Search keeps all thread groups busy during attention score computation and softmax, whereas OpenEvolve restricts this stage to one warp.The comparison identifies unused block capacity in OpenEvolve during score-and-softmax computation.
- MLA Paged Decode: K-Search uses a single block for short sequences, avoiding the reduce pass and extra memory required by fixed larger chunks.
- MLA Paged Decode: K-Search keeps query fragments in registers, loads two chunks ahead, and adapts sequence splits to sequence length.These choices reuse query data, deepen the prefetch pipeline, and adjust splitting instead of applying a fixed strategy.
A.2 Prompt template
The baseline prompt asks LLMs to generate CUDA kernels that match the reference specification while optimizing for the target GPU and fixed tensor characteristics. It also imposes strict output-format and performance-target requirements.
- Generation requirements: The prompt asks for a CUDA kernel implementation optimized for the target GPU while preserving the specification’s computational accuracy.
- Optimization guidance: The prompt directs models to use tensor shapes, dtypes, axes, and GPU memory hierarchy to guide optimization.
- Optimization guidance: The prompt permits third-party libraries but notes that custom implementations can perform better for specialized kernels with known axis constraints.
- Output constraints: Generated code must use XML format with exactly three files having specified names.
- Evaluation: The prompt includes performance targets where lower latency is better.