K-Search：通过协同演化内在世界模型实现LLM内核生成 K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model
用协同演化世界模型替代进化启发式，引导LLM进行GPU内核优化搜索
前置知识
GPU内核（GPU Kernel）
GPU内核是在GPU上并行执行的计算程序，是深度学习系统中训练和推理的核心组件。高效的GPU内核需要精心设计线程块划分（tiling）、内存层次结构利用、同步机制和架构特定指令等。例如FlashAttention用于内存高效的注意力计算，FlashInfer用于高吞吐量的LLM推理服务。内核性能优化面临巨大挑战，因为需要在庞大的设计空间中找到最优配置，且不同硬件架构（如NVIDIA Hopper到Blackwell）的性能权衡完全不同。
本文的核心目标就是自动化生成高性能GPU内核，理解内核优化的复杂性是理解本文动机和方法的基础。
进化搜索（Evolutionary Search）
进化搜索是一种受生物进化启发的优化方法，通过维护一个候选解种群，进行选择、变异和交叉操作来逐步改进解的质量。在代码生成领域，进化搜索结合LLM使用MAP-Elites等质量多样性机制来维护和探索候选程序空间。代表方法包括FunSearch、AlphaEvolve和OpenEvolve。这些方法将LLM视为随机代码生成器，在程序空间中直接搜索，依赖进化启发式驱动搜索进程。
本文的基线方法（OpenEvolve和ShinkaEvolve）都基于进化搜索，理解其局限性才能理解K-Search的创新之处。
世界模型（World Model）
世界模型是一种能够预测环境状态转移的模型，用于模拟行动后果并支持规划和决策。在本文中，世界模型由LLM实例化，用于维护搜索前沿并估计高层优化意图的优先级分数。该模型通过上下文学习不断从执行反馈中更新其对搜索过程的理解，从而动态调整搜索策略。RAP等工作已经证明LLM可以作为世界模型来指导推理和规划任务。
K-Search的核心创新就是将LLM从单纯的代码生成器升级为协同演化的世界模型，这是理解本文方法论的关键概念。
上下文学习（In-Context Learning）
上下文学习是大语言模型的一种能力，能够在不更新模型参数的情况下，通过在输入提示中提供示例来学习新任务或适应新模式。在K-Search中，世界模型通过上下文学习不断吸收执行反馈和搜索历史，从而更新其对搜索过程的理解和信念。这使得系统能够在搜索过程中动态调整优先级分数、提出新的优化假设、以及修剪不成功的搜索分支。
上下文学习是实现世界模型协同演化的核心机制，使LLM能够在不重新训练的情况下持续改进搜索策略。
Split-K并行策略
Split-K是一种GPU内核优化技术，将键值序列分割到多个线程块上并行处理。每个线程块处理序列的一个连续块，计算部分注意力结果并写入临时缓冲区，最后通过规约步骤合并最终结果。这种策略在大批量时能充分利用GPU并行性，但会引入额外的同步开销。在K-Search生成的GQA解码内核中，该策略对长序列和大批次效果显著，但在小批次（如batch_size=1）时反而不如单块方案高效。
这是论文实验部分讨论的一个重要优化技术，展示了K-Search生成的内核的具体优化策略及其权衡。
研究动机
现代机器学习系统对高性能GPU内核的依赖日益增长，但内核优化面临巨大挑战。首先，现代GPU的性能优化需要在庞大的设计空间中探索，涉及线程块划分、内存布局、同步机制和架构特定原语等多个因素的复杂交互。其次，硬件快速迭代增加了额外的复杂性——从NVIDIA Hopper到Blackwell架构的转变引入了新的指令和架构特性，从根本上改变了性能权衡，使得之前优化的内核变得次优。第三，测试一次内核优化可能需要大量的实现工作和反复的试错，编译和分析生成的内核也计算成本高昂，通常要求严格的测试预算限制。现有的LLM进化方法（如OpenEvolve和ShinkaEvolve）将LLM仅仅视为随机代码生成器，依赖MAP-Elites等启发式机制在程序空间中直接选择和变异候选方案。这些方法在处理需要协调的多步结构转换时表现不佳，因为它们缺乏显式的规划能力，经常因低效或错误的中间实现而丢弃有前途的策略。
本文的目标是本文提出Search via Co-Evolving World Model方法，并基于此构建K-Search框架。其核心目标是通过用协同演化的世界模型替代静态搜索启发式，使LLM能够利用其先验领域知识主动引导搜索过程，在优化空间中进行探索。具体而言，该框架旨在显式地将高层算法规划与低层程序实例化解耦，使系统能够在非单调优化路径中导航，同时对临时实现缺陷保持韧性。最终目标是在固定的评估预算内（如120次迭代），为复杂的GPU内核（如GQA、MLA和MoE内核）生成显著优于现有进化搜索方法的高性能实现，甚至超越人类工程师精心优化的基准内核。
与已有工作不同的是，本文的独特切入角度在于将LLM从单纯的代码生成器提升为具有内在规划能力的世界模型。现有进化方法将搜索直接在程序空间中进行，依赖进化启发式驱动进展，而未能充分利用LLM的高层规划和推理能力。K-Search的关键洞见是：有效的搜索需要将算法规划（利用内在和演化的理解）与实现细节分离开来。通过将内核生成形式化为结构化搜索树上的规划问题，由LLM世界模型指导搜索，该方法能够维护搜索前沿、估计优化意图的优先级分数，并通过上下文学习持续从执行反馈中演化其理解。这种范式转变使得系统能够放弃理论上正确但暂时编译错误的策略，而是通过本地精炼机制容忍实现缺陷，从而发现深层结构优化。
核心方法
K-Search将GPU内核生成形式化为一个在结构化搜索树上的规划问题，由一个LLM实例化的世界模型来指导搜索。直觉上，该方法将搜索过程分为两个层次：高层的算法规划（决定应该尝试什么类型的优化）和低层的程序实现（具体实现选定的优化方案）。世界模型负责维护搜索前沿并估计高层优化意图的优先级分数，它通过上下文学习从执行反馈中持续演化其对搜索过程的理解。整个系统通过三个迭代阶段运作：动作选择（从前沿中选择最高优先级的动作）、程序实例化（通过本地精炼生成具体实现）、世界模型更新（分析执行轨迹并更新搜索状态）。这种方法的关键优势在于，它不因暂时的编译错误而丢弃理论上正确的优化策略，而是通过本地精炼机制容忍实现缺陷，让世界模型的理解与内核优化进程共同演化。
K-Search与已有方法的本质区别在于将LLM从被动的代码生成器转变为主动的搜索规划者。现有方法如OpenEvolve和ShinkaEvolve在程序空间中直接搜索，将LLM视为随机采样器，依赖MAP-Elites等进化启发式驱动进展。这些方法本质上将高层优化意图与低层实现耦合在一起，导致一个理论正确的策略可能仅仅因为中间代码的语法错误就被丢弃。K-Search则显式地将搜索空间分为两个层次：高层的动作空间（由优化意图delta描述，如通过填充解决bank冲突）和低层的程序空间（具体实现）。世界模型在高层空间中进行规划，维护搜索树并动态更新优先级分数，而将程序实例化交给专门的本地精炼过程。这种解耦使得系统能够通过Insert、Update和Prune三种树编辑操作来演化搜索状态，从而更有效地导航复杂的非单调优化路径。
方法步骤详情
K-Search的完整工作流程包含三个迭代阶段。第一阶段是动作选择（Action Selection）：从搜索状态 $S_t$ 的当前前沿 $A(S_t)$ 中选择优先级分数 $V$ 最高的动作 $a_t = \arg\max_{a \in A(S_t)} V(a|S_t)$。每个动作 $a = (x_{parent}, \delta)$ 表示一个特定的优化意图 $\delta$ 应用于一个父程序 $x_{parent}$。第二阶段是程序实例化（Program Instantiation）：将选定的计划应用于父实现，通过LLM作为随机策略 $\pi_{code}$ 反复采样具体实现 $x \sim \pi_{code}(\cdot | a_t)$ 并评估 $o = E(x_t)$，直到连续 $K$ 次尝试无改进（停滞条件，实验中 $K=7$）。这确保有效的动作不会因暂时的语法错误而被丢弃。第三阶段是世界模型协同演化（World Model Co-Evolution）：LLM分析执行轨迹，通过三种树编辑操作更新搜索状态——Insert（添加新的子节点扩展当前状态）、Update（根据新证据重新评估现有前沿节点的优先级）和Prune（识别并永久删除不可行或冗余的分支）。搜索状态 $S_t$ 作为显式的搜索树维护，分为Closed节点（已访问，附带最佳程序）和Open节点（待探索的前沿假设）。
技术新颖性
K-Search的技术新颖性体现在多个层面。首先，将LLM作为协同演化的世界模型来指导代码搜索是一个全新的范式，不同于将LLM视为被动代码生成器的现有方法。世界模型表示搜索状态转移分布 $P_{model}(S_{t+1}|S_t, a_t)$，并通过上下文学习持续从执行反馈中演化其理解。其次，显式的搜索树结构使得系统能够进行有意义的树编辑操作（Insert、Update、Prune），而不是仅在程序空间中盲目变异。第三，本地精炼机制（stagnation limit $K=7$）将逻辑推理与实现细节隔离，允许系统对临时实现缺陷保持韧性。第四，搜索空间的层次化分解——高层优化意图空间和低层程序空间——使得系统能够探索非单调优化路径，其中中间步骤可能不会立即带来性能提升。实验表明，这种范式在MLA Paged Decode内核的案例研究中展示了高效的搜索过程：从初始假设（fused_multi_head）出发，通过树演化逐步深入，最终发现chunk32_prescale_vectorized这一全局最优解。
实验结果
K-Search在四个代表性的FlashInfer内核上进行了全面评估。在所有内核上，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 prefill内核上，K-Search得分57.4，相比OpenEvolve的19.5和ShinkaEvolve的11.3分别实现了2.95倍和5.10倍改进。GQA decode内核上，K-Search达到76.0分，相比OpenEvolve（44.2）和ShinkaEvolve（27.7）分别提升1.72倍和2.74倍。MLA decode内核上，K-Search维持47.1分，相比基线提升18%和36%。在GPUMODE TriMul竞赛中，K-Search以1030微秒的几何平均延迟达到最优性能，超越了包括使用25600次迭代的RL加进化方法（TTT，1161微秒）在内的所有现有方案。值得注意的是，K-Search在GQA decode的小批次工作负载上略逊于基线，这是因为其采用的split-K并行策略在小批次时引入了不必要的同步开销。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GQA Paged Decode | Score（速度提升百分比） | 76.0 | OpenEvolve 44.2 / ShinkaEvolve 27.7 | 1.72x / 2.74x |
| MLA Paged Decode | Score（速度提升百分比） | 47.1 | OpenEvolve 39.9 / ShinkaEvolve 34.7 | 18% / 36% |
| MLA Paged Prefill | Score（速度提升百分比） | 57.4 | OpenEvolve 19.5 / ShinkaEvolve 11.3 | 2.95x / 5.10x |
| FP8 MoE | Score（速度提升百分比） | 44.1 | OpenEvolve 3.09 / ShinkaEvolve 27.9 | 14.3x / 1.58x |
| GPUMODE TriMul (H100) | Latency (微秒, lower is better) | 1030 | TTT-Discover 1161 / Zeyu Shen 1140 / shiyegao 1074 | SOTA，优于所有现有方案 |
局限与改进
论文承认了若干局限性。首先，在GQA decode的小批次工作负载（batch_size=1和batch_size=16）上，K-Search生成的内核性能不如OpenEvolve和ShinkaEvolve，这是因为其采用的split-K并行策略在小批次时引入了不必要的同步开销，而基线方法的单块设计在无协调需求时更高效。其次，实验设置的评估预算限制为120次迭代，虽然这在资源效率方面是合理的，但可能不足以充分探索某些极其复杂的优化空间。第三，世界模型的演化目前仅通过上下文学习实现，没有涉及模型微调或强化学习，这可能限制了其在更大搜索空间中的学习能力。第四，论文主要关注FlashInfer相关的注意力和MoE内核，对于其他类型的GPU内核（如卷积、矩阵乘法等）的泛化能力尚未验证。此外，K-Search生成的内核虽然显著优于进化基线，但很少能超越人类专家精心优化的FlashInfer内核，表明自动化内核生成仍有较大提升空间。
独立分析的弱点
K-Search存在几个值得深入分析的弱点。第一，本地精炼阶段的停滞阈值 $K=7$ 是固定的，不同内核可能需要不同的精炼深度；对于Triton内核，论文将 $K$ 降至5，表明这个超参数需要针对不同场景调优，未来可以考虑自适应停滞检测机制。第二，世界模型的演化完全依赖上下文学习，没有利用模型微调或强化学习来积累长期知识，这限制了跨任务的知识迁移。第三，搜索树的规模随着迭代增长，世界模型需要处理越来越长的上下文，可能导致决策质量下降或计算成本增加，论文未讨论搜索树的剪枝策略对大规模搜索的影响。第四，实验仅使用了120次评估预算，对于实际工业应用可能不够；同时，实验在固定硬件（H100和B200）上进行，对于跨架构泛化能力的评估不足。改进方向包括引入自适应停滞检测、结合轻量级微调机制、以及开发更高效的搜索树管理策略。
未来方向
论文作者和基于本文成果可以延伸出多个未来研究方向。首先，世界模型可以从纯上下文学习扩展到结合强化学习或轻量级微调，使系统能够在任务间积累和迁移优化知识。其次，搜索空间可以扩展到更多的内核类型和硬件架构，包括Blackwell架构的新指令优化、多GPU内核、以及非注意力类型的算子。第三，可以探索将K-Search的规划框架与其他代码生成技术（如程序合成、形式化验证）结合，提高生成内核的正确性保证。第四，世界模型的可解释性研究可以帮助理解LLM如何学习GPU内核优化的领域知识，以及如何更好地利用这种知识。第五，将该框架应用于更大规模的搜索预算和更复杂的优化任务，如整模型推理优化或自动编译器调优，是具有潜力的研究方向。最后，探索多世界模型协作或层次化世界模型架构，可能进一步提升复杂优化任务的搜索效率。
复现评估
论文在复现性方面提供了较好的支持。代码已在GitHub开源（https://github.com/caoshiyi/K-Search），包含完整的K-Search框架实现。实验使用了标准化的评估环境：NVIDIA H100和B200 GPU，CUDA 12.8，FlashInfer 0.5.3，PyTorch 2.8.0。论文使用了FlashInfer-Bench作为评估器，所有候选实现必须是CUDA代码并通过功能正确性测试。评估使用了固定的工作负载集合（来自真实流量捕获的152个总工作负载轨迹），确保了公平比较。论文重复每个方法三次并报告均值曲线和最小-最大范围。对于基线方法，使用了gemini-3-pro-preview作为LLM，并遵循默认配置。然而，复现面临的主要挑战包括：需要昂贵的GPU硬件（H100/B200）、较长的搜索时间（120次迭代每次需要编译和评估）、以及对FlashInfer-Bench环境的依赖。此外，GPUMODE TriMul实验使用了GPT-5.2和Gemini-3-Pro的组合，这些商业API的成本和可用性可能影响复现。
论文图表