UC Berkeley 2026年3月5日 15:00
关键词： GPU 内核优化 、启发式、 世界模型 、 演化搜索 、代码生成
,25分钟
在深度学习框架的底层，GPU 内核（kernel）的优化是一场永无止境的军备竞赛。从 FlashAttention 到 FlashInfer，每一个算子库的背后都凝聚着专家们对 GPU 架构的极致理解。
然而，随着硬件快速迭代（如从 NVIDIA Hopper 到 Blackwell）和模型结构多样化（如 Mamba、MoE），手动优化内核的成本已经高到难以承受。于是，一个自然的想法是：能否让大语言模型（LLM）自动生成高性能的 GPU 内核？
已有的方法大多将 LLM 视为随机代码生成器，结合遗传算法或质量多样性搜索（如 OpenEvolve、ShinkaEvolve）在程序空间中进行演化。 然而，这些方法往往缺乏高层规划能力： 一个理论上优秀的优化策略可能因为中间步骤的编译错误而被过早丢弃 ，而需要多步协同的结构性变换（如先调整内存布局再向量化）更是难以实现。
问题出在哪里？
核心在于 将 LLM 仅用作代码生成器，而非规划引擎 。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
- K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model
- https://arxiv.org/pdf/2602.19128v1
- 代码：https://github.com/caoshiyi/K-Search
,25分钟
相关推荐
- CUDA Kernel 自动生成新范式！多轮RL+反馈迭代方案 Kevin 让生成速度超o4-mini 41%、正确性领先44%
- Meta 提出 TritorX：面向 ML ASIC 的 Agentic 算子生成系统，84.7% OpInfo 覆盖与 2 万+测试验证，简化加速器软件生态构建
- 跨平台 Kernel 自动生成综述：基于 LLM 的监督微调+强化学习，适配 NVIDIA/AMD/华为 NPU
本文要介绍的 K-Search 方法来自 UC Berkeley 团队提出了的一种全新范式，文章名是《K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model》，作者将 LLM 重新定位为 共同演化的世界模型，在搜索过程中动态维护一个显式的规划树 ，从而实现对 GPU 内核优化空间的深度探索。
实验表明， K-Search 在 GQA、MLA、MoE 等复杂内核上平均性能达到基线方法的 2.1 倍 ，在 MoE 任务上最高提升 14.3 倍 ，并在 GPUMODE TriMul 竞赛中登顶 H100 榜首。下面，我们将深入解读这一工作的核心创新与技术细节。
本文目录
- 一、相关工作：从自动调优到 LLM 演化
- 二、核心创新：共同演化的世界模型
- 2.1 问题形式化
- 2.2 搜索状态与行动
- 2.3 工作流程
- 四、案例研究：MLA Paged Decode 中的世界演化
- 初始化（Round 1）
- 中期演化（Rounds 14-34）
- 后期收敛（Round 95）
- 五、实验验证：全面超越现有方法
- 5.1 主要结果
- 5.2 细粒度分析
- 5.3 深入解读：为何 K-Search 更优？
- 六、扩展验证：GPUMODE TriMul 登顶
- 结论与展望
交流加群请在 NeuralTalk 公众号后台回复：加群
一、相关工作：从自动调优到 LLM 演化
在进入 K-Search 之前，我们先回顾一下 GPU 内核优化的技术脉络。
传统编译器自动调优 ：以 TVM 及其 Ansor 调度器为代表，通过成本模型在庞大的调度空间中搜索最优配置。这类方法擅长处理张量程序的结构化变换， 但往往受限于预定义的调度原语，难以引入全新的算法级优化。
领域特定语言（DSL） ：Triton 和 CuTe 等工具让开发者可以更高层次地表达计算逻辑，同时保留对内存层次和并行的精细控制。然而， DSL 仍然依赖专家经验，且无法自动探索非局部的优化组合。
LLM+演化搜索 ：近期，FunSearch、OpenEvolve 等将 LLM 与遗传算法结合，用 LLM 生成变体，用进化启发式（如 MAP-Elites）选择候选。 这类方法直接在程序空间操作，但存在两个根本局限：
- 缺乏规划 ——LLM 只看到当前程序片段，无法规划需要多步才能见效的优化路径；
- 对噪声敏感 ——一次编译错误或性能回退就可能导致整个分支被抛弃。
K-Search 的洞察在于：LLM 不仅知道如何写代码，更内化了大量关于 GPU 架构和优化策略的 先验知识 。
如果能将这些知识组织成一个可演化的“世界模型”，让 LLM 在搜索过程中不断修正自己的理解 ，就能从单纯的代码生成跃升为 自主规划 。
二、核心创新：共同演化的世界模型
K-Search 的核心思想是将 Kernel 生成视为一个 规划问题 ，搜索空间由一棵显式的 搜索树 表示，而树的维护与扩展由一个基于 LLM 的 世界模型 负责。
这个世界模型有两个关键特性：
- 内蕴先验 ：LLM 本身具备关于 GPU 优化（如合并访问、bank 冲突避免、张量核心使用）的丰富知识，可以初始估计不同优化方向的潜力。
- 共同演化 ：世界模型会随着搜索进程不断更新——根据实际执行反馈（编译日志、性能数据）调整对各个优化动作的置信度，插入新的优化假设，剪除无效分支。这种动态更新完全通过 上下文学习 实现，无需微调。
2.1 问题形式化
给定一个内核程序 ，评估器 返回观测元组 ，其中 表示正确性， 是延迟， 是元数据如编译器输出、Profiler 日志问题。目标是最大化相对于参考基线的加速比 ：
是 SoTA 程序实现，在固定的评估预算 B 内，找到最优程序 。
2.2 搜索状态与行动
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 1 | K-Search 框架概述。该框架基于搜索状态 运行， 的结构为搜索树。这棵树由封闭节点（蓝色，已访问状态，带有类似 的程序）和开放节点前沿（橙色，待处理假设，如u13）组成。工作流程分为三个阶段循环进行：（1）动作选择，根据世界模型估计的优先级分数V从前沿中检索最有前景的动作节点；（2）局部优化，随机策略 采样具体实现，直至停滞；（3）世界模型更新，大语言模型（LLM）对轨迹进行推理，通过插入（添加新动作）、更新（调整V，例如u11从0.9降至0.6）和修剪（移除不太有前景的节点，如u10）来更新搜索树。三个阶段形成闭环优化逻辑，先依据世界模型的优先级得分选最优动作，再通过局部细化生成具体程序并验证，最后根据执行反馈更新搜索树，实现规划与实现解耦，规避临时实现缺陷导致的有效策略丢失。
搜索状态 被组织为一棵树，包含两类节点如图 1：
- CLOSED 节点 （蓝色）：已访问过的程序实例，附有最佳实现（如 x₁₂）。
- OPEN 节点 （橙色虚线）：待探索的 行动 ，每个行动是一个元组(, δ)，表示对父程序 应用某种优化意图 δ（例如“通过 padding 解决 bank 冲突”）。每个 OPEN 节点还有一个由世界模型动态估计的 优先级分数 。
世界模型通过三种树编辑操作演化状态：
- 插入（Insert） ：根据当前最佳程序的表现，衍生出新的子节点（优化意图）。
- 更新（Update） ：根据执行反馈，调整现有 OPEN 节点的优先级（如图中 u₁₁ 从 0.9 降为 0.6）。
- 剪枝（Prune） ：移除被证实无效或冗余的分支（如 u₁₀）。
这种树结构将 高层规划（选择优化意图）与低层实现（生成具体代码）解耦 ，使得 系统可以坚持一个有前景的优化方向，即使中间实现有缺陷也能通过局部重试修复。
2.3 工作流程
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 1 | K-Search 框架概览。搜索状态由 CLOSED 节点（蓝色，已访问程序）和 OPEN 节点（橙色，待探索行动）构成。世界模型通过插入、更新、剪枝操作演化树结构，优先级 V 动态调整。
K-Search 的迭代流程如图 1 所示，算法 1 放在本段后面。算法主要包含三个阶段：
- 行动选择 ：从当前前沿（OPEN 节点）中选出优先级最高的行动 即 。
- 程序实例化（局部细化） ：将 LLM 作为随机策略 ，反复采样具体实现 ，并评估，直到连续 K 次无改进（K=7）。这一阶段保证一个优化意图被充分探索，避免因语法错误或小 bug 而放弃。
- 世界模型演化 ：将最佳实现 及其观测 反馈给世界模型，后者根据新证据更新搜索树（插入新节点、调整优先级、剪枝）。 ’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
算法 1 K-Search。该算法将 K-Search 三阶段核心逻辑程序化，以任务规格 、预算 为输入，先通过 选取最优动作，再在局部细化循环中以 采样生成候选程序 ，经评估器 验证后，通过 更新搜索状态，以停滞阈值 约束采样次数，在有限 内高效筛选出最优内核 ，实现规划与实现解耦，大幅降低无效探索成本。
这种“选择-细化-演化”的循环，使得搜索既能聚焦于高潜力的方向，又能动态适应新信息，实现非单调的优化路径探索。
四、案例研究：MLA Paged Decode 中的世界演化
为了更直观地理解世界模型如何演化，MLA 是一种多级注意力机制，在解码阶段需要高效处理分页 KV 缓存。我们以 MLA Paged Decode 内核的优化过程为例，如图 2。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 2 | K-Search 搜索轨迹可视化。它追踪了在 MLA 分页解码内核上各搜索轮次中搜索状态的演变。一轮（round）对应一次候选程序评估。节点代表动作（蓝色=已关闭，橙色=已打开），并标注了其实例化程序的性能（已关闭节点）或优先级分数（已打开节点）。时间线突出显示了内核是如何改进的，以及大型语言模型（LLM）如何基于不断深化的理解动态插入新假设、更新信念并修剪前景较差的分支。轨迹清晰呈现了从初始假设到全局最优的演进：通过剪枝独立头分支、重构 split-K 策略，最终通过 chunk32 预缩放矢量化实现性能突破，直观体现了协同进化模型对优化策略的动态校准能力。
上图是 MLA Paged Decode 内核的世界模型演化示例。随着搜索进行，模型插入新节点（如 register_resident_scaling）、更新优先级（u₁₁ 从 0.9 降至 0.6，结合图 1）、剪枝无效分支（u₁₀，结合图 1），最终收敛到高性能实现。
下面我们逐步来看。
初始化（Round 1）
世界模型在根节点下插入三个初始行动：fused_multi_head（融合多头）、split_k decoding（Split-K 解码）、independent_heads（独立头）。
基于先验知识，模型认为 fused_multi_head 能通过共享 CKV 头减少 16 倍全局内存流量，赋予其最高优先级 V，并选择它进行细化。
中期演化（Rounds 14-34）
在细化 fused_multi_head 的过程中，程序得分 提升到 34。世界模型根据这一反馈：
- 插入 新节点：如 register_resident_scaling（寄存器驻留缩放）、occupancy_tuned_chunk32（占用率优化的块大小），以深化成功分支。
- 更新 优先级：将 sibling 节点 independent_heads 的优先级下调，因为融合头的成功意味着独立处理不再有竞争力。
- 剪枝 ：移除明显无效的分支（图 1 中 u₁₀）。
后期收敛（Round 95）
当 fused_multi_head 分支性能趋于平稳，世界模型将 注意力转向 之前优先级较低的 split_k decoding 分支。经过细化，发现其更适合某些批处理场景，于是将其 优先级提高 ，继续探索。
这个案例生动展示了世界模型如何根据实际证据动态调整搜索策略：不是盲目跟随初始假设，而是 不断修正信念，平衡探索与利用。
五、实验验证：全面超越现有方法
K-Search 在四个代表性的复杂内核上与 OpenEvolve 和 ShinkaEvolve 进行了对比：GQA 解码、MLA 解码、MLA 预填充、MoE。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
表 2 | 用于评估的代表性内核，包含内核名称、适配架构、对应的模型场景，以及各内核的特性与优化挑战，因篇幅省略了头数等详细配置，优化挑战基于 FlashInfer 的实现设定。所选内核均为 LLM 推理服务中的核心负载，覆盖 Hopper、Blackwell 两大主流 GPU 架构，各内核面临带宽、延迟、内存、数据路由等不同类型优化难题，能全面检验 K-Search 在多样场景下的优化能力。
所有实验均使用 NVIDIA H100/B200 GPU，评估预算 120 次，每次评估包括编译和性能测试。表 2 总结了内核的关键特征。
5.1 主要结果
图 3a 展示了三种方法在 120 次迭代中的最佳得分曲线。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 3 | 主要结果（3 次重复实验）。(a) 对比三大系统在 120 次迭代中各内核的最优累计得分；(b) 对所有对比系统做单工作负载性能分析；(c) 展示各系统最优内核相对 FlashInfer 基线实现指定加速比的工作负载占比。三组子图从迭代趋势、单负载性能、加速比覆盖率三个维度验证 K-Search 优势，其得分增长更快、多数负载性能更优，高加速比覆盖率远超 OpenEvolve 和 ShinkaEvolve，尤其在复杂 MoE 内核上性能提升尤为显著。
K-Search 在所有内核上均显著领先 ，最终平均得分 56.13，是 OpenEvolve（26.68）的 2.10 倍，ShinkaEvolve（25.37）的 2.21 倍。
- 最令人印象深刻的当属 MoE 内核：K-Search 得分 44.1，相比 OpenEvolve 的 3.09 提升 14.3 倍，相比 ShinkaEvolve 的 27.9 提升 1.58 倍。
- 对于 MLA 预填充，K-Search 得分 57.4，分别是 OpenEvolve（19.5）和 ShinkaEvolve（11.3）的 2.95 倍和 5.10 倍。
5.2 细粒度分析
图 3b 给出了每个工作负载实例上的性能散点图。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 3 | 主要结果（3 次重复实验）。(a) 对比三大系统在 120 次迭代中各内核的最优累计得分；(b) 对所有对比系统做单工作负载性能分析；(c) 展示各系统最优内核相对 FlashInfer 基线实现指定加速比的工作负载占比。
在总共 152 个测试轨迹中，K-Search 在绝大多数负载上优于基线。 但有趣的是，在 GQA 解码的小批量（batch size=1 或 16）任务上，K-Search 略逊于 OpenEvolve。 原因是 K-Search 采用了 Split-K 策略（将序列分块并行），对于小批量而言同步开销过大；而基线简单的单块设计反而更优 。这恰恰体现了 K-Search 的潜力——如果能进一步引入自适应策略，根据批量大小动态选择并行方式，可能覆盖更广的场景。
图 3c 展示了达到不同加速比阈值的负载比例。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
图 3 | 主要结果（3 次重复实验）。(a) 对比三大系统在 120 次迭代中各内核的最优累计得分；(b) 对所有对比系统做单工作负载性能分析；(c) 展示各系统最优内核相对 FlashInfer 基线实现指定加速比的工作负载占比。
- 在 GQA 解码上，K-Search 在 100%负载上达到 ≥0.36 加速比，而 ShinkaEvolve 仅覆盖 50%；在 ≥0.50 加速比上，K-Search 覆盖 87.5%，OpenEvolve 和 ShinkaEvolve 分别只有 50.0%和 39.6%。
- 对于 MLA 预填充，K-Search 在 57.9%负载上达到 ≥0.40 加速比，而基线完全无法触及这一阈值。
5.3 深入解读：为何 K-Search 更优？
FP8 MoE
以 FP8 MoE 内核为例。MoE 需要为每个 token 从 256 个专家中选出 top-8，然后执行门控、上投影、SiLU、下投影。
- K-Search 的实现采用每 token 一个线程块（256 线程），通过 warp 级协作（__shfl_down_sync）并行选出 top-8，避免串行。
- 而 OpenEvolve 使用持久化内核，每个块通过原子操作从全局计数器获取下一个任务，导致巨大开销；
- ShinkaEvolve 则让每个线程循环扫描所有专家分数，性能低下。
在专家 FFN 计算阶段，K-Search 利用张量核心（WMMA）和双缓冲，将数据加载与计算重叠，并跳过零 tokens 的专家。这些精妙设计正是世界模型引导下逐步细化出来的。
GQA 解码
再如 GQA 解码，K-Search 采用 Split-K 策略将长序列分块并行，并使用双缓冲（甚至三缓冲）隐藏访存延迟，而基线单块串行处理序列，无法有效利用 GPU 的并行能力。
六、扩展验证：GPUMODE TriMul 登顶
除了 FlashInfer 内核，K-Search 还在 GPUMODE 平台的 TriMul 任务上进行了测试。TriMul 是 AlphaFold3 中的核心模块，涉及 4 维张量的复杂运算。下表 3 列出了前几名提交的成绩。
’ fill=‘%23FFFFFF’%3E%3Crect x=‘249’ y=‘126’ width=‘1’ height=‘1’%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
表 3 | NVIDIA H100 上的顶级 GPUMODE TriMul 提交结果。该表格展示 NVIDIA H100 显卡上 GPUMODE 平台 TriMul 任务的顶级提交结果，包含提交 ID、开发语言、使用模型、迭代次数，以及各方案在固定基准案例下的几何平均延迟，延迟数值越小性能越好。TriMul 是蛋白质结构预测的核心模块，K-Search 仅用 300 次迭代，结合 GPT-5.2 与 Gemini-3-Pro 就实现 1030μs 的最低延迟，远少于竞品迭代次数，印证其在科学计算内核优化中的高效性。
K-Search 在 300 次迭代内（前 150 次用 GPT-5.2，后 150 次用 Gemini-3-Pro），无需任何种子程序，便达到了 1030μs 的几何平均延迟， 超越此前所有人工设计和自动搜索方案 ，包括结合强化学习的 TTT-Discover（1161μs）。
结论与展望
K-Search 通过将 LLM 重塑为共同演化的世界模型，实现了对 GPU 内核优化空间的高效规划式搜索。其核心贡献在于： 将高层优化意图与低层代码实现解耦，并让世界模型动态更新对搜索方向的信念 。实验证明， 这种范式在复杂内核上显著超越现有演化方法，甚至在【某些任务】上【超越】专家手工优化。
未来的工作可以沿着多个方向延伸：
- 将世界模型的 更新从上下文学习升级为在线微调 ；
- 将 树搜索与强化学习结合 ，实现更精细的优先级估计；
- 将方法 推广到更多硬件 （如 AMD、定制加速器）和 更广泛的程序优化任务 。
可以预见，LLM 作为世界模型的潜力远未被充分挖掘，K-Search 只是迈出了第一步。
相关推荐
- NVIDIA 技术博客：基于 DeepSeek-R1 与推理计算扩展的 GPU Kernel 自动生成
- NVIDIA, UCB提出Kernel Blaster：记忆增强上下文 RL 让 CUDA 内核几何平均1.43×-2.50×加速！
- MultiKernelBench：首个覆盖英伟达GPU、华为NPU、谷歌TPU的Kernel生成基准，涵盖14类算子285项测试
交流加群请在 NeuralTalk 公众号后台回复：加群
GPU · 目录
继续滑动看下一个
NeuralTalk
向上滑动看下一个