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

arXiv: https://arxiv.org/abs/2602.19128 · HTML: https://mlsanity.com/paper/2602.19128/

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 31 of its 43 references that are in the map.
- First posted on arXiv 22 Feb 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 30 of 31)

- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [AlphaEvolve: A coding agent for scientific and algorithmic discovery](https://mlsanity.com/paper/2506.13131/) (2025)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Illuminating search spaces by mapping elites](https://mlsanity.com/paper/1504.04909/) (2015)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Large Language Models as Commonsense Knowledge for Large-Scale Task Planning](https://mlsanity.com/paper/2305.14078/) (2023)
- [KernelBench: Can LLMs Write Efficient GPU Kernels?](https://mlsanity.com/paper/2502.10517/) (2025)
- [Leveraging Pre-trained Large Language Models to Construct and Utilize World Models for Model-based Task Planning](https://mlsanity.com/paper/2305.14909/) (2023)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [ShinkaEvolve: Towards Open-Ended And Sample-Efficient Program Evolution](https://mlsanity.com/paper/2509.19349/) (2025)
- [Is Your LLM Secretly a World Model of the Internet? Model-Based Planning for Web Agents](https://mlsanity.com/paper/2411.06559/) (2024)
- [Learning to Discover at Test Time](https://mlsanity.com/paper/2601.16175/) (2026)
- [Kevin: Multi-Turn RL for Generating CUDA Kernels](https://mlsanity.com/paper/2507.11948/) (2025)
- [WebEvolver: Enhancing Web Agent Self-Improvement with Coevolving World Model](https://mlsanity.com/paper/2504.21024/) (2025)
- [CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning](https://mlsanity.com/paper/2507.14111/) (2025)
- [Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization](https://mlsanity.com/paper/2509.14279/) (2025)
- [Astra: A Multi-Agent System for GPU Kernel Performance Optimization](https://mlsanity.com/paper/2509.07506/) (2025)
- [CudaForge: An Agent Framework with Hardware Feedback for CUDA Kernel Optimization](https://mlsanity.com/paper/2511.01884/) (2025)
- [AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs](https://mlsanity.com/paper/2507.05687/) (2025)
- [FlashInfer-Bench: Building the Virtuous Cycle for AI-driven LLM Systems](https://mlsanity.com/paper/2601.00227/) (2026)
- [TritonForge: Profiling-Guided Framework for Automated Triton Kernel Optimization](https://mlsanity.com/paper/2512.09196/) (2025)
- [EvoEngineer: Mastering Automated CUDA Kernel Code Evolution with Large Language Models](https://mlsanity.com/paper/2510.03760/) (2025)

## Built on by (most cited first, 24 of 24)

- [ASPIRE: Agentic /Skills Discovery for Robotics](https://mlsanity.com/paper/2607.00272/) (2026)
- [Structured Scaling of AI Discovery Across Diverse Scientific Domains](https://mlsanity.com/paper/2604.19341/) (2026)
- [Dream-RSI: Recursive Self-Improvement through Evolving Worlds](https://mlsanity.com/paper/2609.14858/) (2026)
- [KernelBench-Verified: Do LLM-Generated Kernels Actually Beat PyTorch?](https://mlsanity.com/paper/2607.16241/) (2026)
- [Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization](https://mlsanity.com/paper/2603.28342/) (2026)
- [CuTeGen: An LLM-Based Agentic Framework for Generation and Optimization of High-Performance GPU Kernels using CuTe](https://mlsanity.com/paper/2604.01489/) (2026)
- [Improving Efficiency of GPU Kernel Optimization Agents using a Domain-Specific Language and Speed-of-Light Guidance](https://mlsanity.com/paper/2603.29010/) (2026)
- [CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution](https://mlsanity.com/paper/2608.12629/) (2026)
- [FrontierSmith: Synthesizing Open-Ended Coding Problems at Scale](https://mlsanity.com/paper/2605.14445/) (2026)
- [What Do Evolutionary Coding Agents Evolve?](https://mlsanity.com/paper/2605.20086/) (2026)
- [JAXBench: Benchmarking Autonomous TPU Kernel Optimization](https://mlsanity.com/paper/2607.20466/) (2026)
- [SLAI T-Rex: Full-Parameter Post-training of the DeepSeek-V4 Family on Ascend SuperPOD](https://mlsanity.com/paper/2607.20145/) (2026)
- [Humanize: Judgement Engineering for Agentic Coding](https://mlsanity.com/paper/2610.08900/) (2026)
- [Automated Membership Inference Attacks (AutoMIA): Discovering MIA Signal Computations using LLM Agents](https://mlsanity.com/paper/2603.19375/) (2026)
- [Ave: Guiding Agentic GPU Optimization Using Data-Flow Invariants](https://mlsanity.com/paper/2604.18616/) (2026)
- [Learning When to Optimize: Verified Optimization Skills from Expert GPU-Kernel Lineages](https://mlsanity.com/paper/2605.28213/) (2026)
- [GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization](https://mlsanity.com/paper/2605.31464/) (2026)
- [PTXBench: Benchmarking and Adapting LLMs for GPU Kernel Optimization with Architecture-specific PTX](https://mlsanity.com/paper/2608.17379/) (2026)
- [KernelOPT: Dispatch-Aware Agentic Search for GPU Kernel Optimization](https://mlsanity.com/paper/2609.30059/) (2026)
- [KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation](https://mlsanity.com/paper/2609.33074/) (2026)
- [GPUPhysBench: Benchmarking Coding Agents for Correct and Efficient GPU Physics Simulation](https://mlsanity.com/paper/2609.35639/) (2026)
- [Evolving in Thought Space: Training a Small Model at Test Time Unlocks Better Discoveries](https://mlsanity.com/paper/2610.06269/) (2026)
- [Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling](https://mlsanity.com/paper/2610.07742/) (2026)
- [JANUS: Enhancing Long-Horizon Autonomous Research with Co-Evolving World Models](https://mlsanity.com/paper/2610.13194/) (2026)

## Abstract (arXiv)

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.