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

📄 arXiv: 2602.19128 📥 PDF

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

分类: cs.AI

发布日期: 2026-02-28

💡 一句话要点

K-Search：通过协同演化的内在世界模型生成LLM Kernel，显著提升GPU Kernel优化效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: GPU Kernel优化 LLM 进化搜索 世界模型 代码生成

📋 核心要点

  1. 现有GPU Kernel优化方法依赖启发式进化搜索，将LLM视为随机代码生成器，缺乏规划能力，难以处理复杂Kernel优化。
  2. K-Search提出协同演化的世界模型，利用LLM的先验知识指导搜索，解耦高层规划和底层实现，提升优化效率。
  3. 实验表明，K-Search在FlashInfer的GQA、MLA、MoE Kernel上显著优于现有方法，最高提升14.3倍，并在GPUMode TriMul任务中达到SOTA。

📝 摘要（中文）

优化GPU Kernel对于高效的现代机器学习系统至关重要，但由于设计因素的复杂相互作用和硬件的快速发展，这仍然具有挑战性。现有的自动化方法通常仅仅将大型语言模型（LLM）视为启发式引导的进化循环中的随机代码生成器。这些方法在需要协调的多步骤结构转换的复杂Kernel上表现不佳，因为它们缺乏明确的规划能力，并且由于效率低下或不正确的中间实现而经常丢弃有希望的策略。为了解决这个问题，我们提出了通过协同演化世界模型进行搜索的方法，并在此基础上构建了K-Search。通过用协同演化的世界模型取代静态搜索启发式方法，我们的框架利用LLM的先验领域知识来指导搜索，主动探索优化空间。这种方法将高层算法规划与低层程序实例化显式解耦，使系统能够导航非单调优化路径，同时保持对临时实现缺陷的弹性。我们在FlashInfer中各种复杂的Kernel上评估了K-Search，包括GQA、MLA和MoE Kernel。结果表明，K-Search显著优于最先进的进化搜索方法，平均提高了2.10倍，在复杂的MoE Kernel上提高了高达14.3倍。在GPUMode TriMul任务中，K-Search在H100上实现了最先进的性能，达到1030us，超过了之前的进化和人工设计的解决方案。

🔬 方法详解

问题定义：论文旨在解决GPU Kernel优化问题，特别是针对复杂Kernel，现有方法如基于LLM的进化搜索，将LLM作为随机代码生成器，缺乏有效的规划和搜索策略，容易陷入局部最优，难以处理需要多步骤结构转换的优化任务，且对中间实现缺陷敏感。

核心思路：论文的核心思路是引入一个协同演化的世界模型，该模型与Kernel代码共同进化。世界模型利用LLM的先验知识，对优化过程进行规划和指导，从而避免盲目搜索。通过解耦高层算法规划和底层程序实例化，系统可以探索非单调的优化路径，并容忍临时的实现缺陷。

技术框架：K-Search框架包含两个主要部分：Kernel代码和世界模型。Kernel代码是待优化的GPU Kernel实现。世界模型是一个LLM，它维护着对Kernel优化过程的理解，并根据当前状态提出优化策略。框架通过一个迭代过程工作：1) 世界模型根据当前Kernel代码的状态，提出一个或多个优化方案；2) 这些方案被实例化为新的Kernel代码；3) 新的Kernel代码被评估，其性能反馈给世界模型；4) 世界模型根据反馈更新其对优化空间的理解，并提出新的优化方案。这个过程不断迭代，直到找到最优的Kernel代码。

关键创新：最重要的技术创新点在于协同演化的世界模型。与传统的进化搜索方法不同，K-Search不是简单地随机生成和评估代码，而是利用LLM的先验知识来指导搜索过程。世界模型能够理解Kernel代码的结构和性能瓶颈，并提出有针对性的优化方案。此外，世界模型还能够学习优化过程中的经验，从而不断提高其规划能力。

关键设计：K-Search的关键设计包括：1) 世界模型的选择：论文使用LLM作为世界模型，因为LLM具有强大的代码理解和生成能力；2) 优化方案的表示：优化方案被表示为对Kernel代码的修改指令，这些指令可以被LLM理解和执行；3) 性能反馈机制：Kernel代码的性能通过运行时间和资源利用率等指标来衡量，这些指标被反馈给世界模型，用于更新其对优化空间的理解；4) 探索-利用策略：K-Search采用一种探索-利用策略，在探索新的优化方案和利用已知的有效方案之间进行平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

K-Search在FlashInfer的GQA、MLA和MoE Kernel上进行了评估，结果表明其显著优于最先进的进化搜索方法，平均提高了2.10倍，在复杂的MoE Kernel上提高了高达14.3倍。在GPUMode TriMul任务中，K-Search在H100上实现了最先进的性能，达到1030us，超过了之前的进化和人工设计的解决方案，证明了其在复杂Kernel优化方面的强大能力。

🎯 应用场景

K-Search具有广泛的应用前景，可用于自动优化各种GPU Kernel，包括深度学习、科学计算、图形渲染等领域。通过降低Kernel优化的门槛，K-Search可以帮助开发者更高效地利用GPU资源，加速应用程序的开发和部署。未来，K-Search有望成为一种通用的GPU Kernel优化工具，推动高性能计算的发展。

📄 摘要（原文）

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.