AdaExplore：基于失败驱动的自适应与多样性保留搜索的高效内核生成

# AdaExplore：面向高效 Kernel 生成的失败驱动适配与多样性保持搜索
Source: https://arxiv.org/html/2604.16625 Weihua Du1, Jingming Zhuo2, Yixin Dong1, Andre He1, Weiwei Sun1, Zeyu Zheng1, Manupa Karunaratne3, Ivan Fox3, Tim Dettmers1, Tianqi Chen1, Yiming Yang1, Sean Welleck1 
1卡内基梅隆大学，2华盛顿大学，3Arm有限公司。 {weihuad, swelleck}@cs.cmu.edu

###### Abstract
近期的大语言模型（LLM）智能体在利用执行反馈进行测试时自适应方面展现出巨大潜力。然而，稳健的自我改进远未解决：大多数方法仍将每个问题实例孤立对待，未能积累可复用的知识。这一局限性在 Triton 等域特定语言中尤为明显，它们在 LLM 预训练数据中占比偏低。其严格的约束条件和非线性的优化地形，使得直接生成和本地细化变得不可靠。为此，我们提出了 AdaExplore，这是一个基于智能体的框架，专为对性能至关重要的内核（Kernel）代码生成提供自我改进能力。该框架通过两个互补阶段实现：基于失败的自适应（failure-driven adaptation）与多样性保持搜索（diversity-preserving search）。两者协同工作，在不依赖额外微调或外部知识的情况下，共同提升代码的正确性与优化性能。在自适应阶段，智能体合成训练任务，并将反复出现的失败转化为可复用的有效性规则记忆库，帮助后续生成的代码保持在可行集内。在搜索阶段，智能体将候选内核组织为树状结构，交替执行小幅度的本地细化和较大规模的结构重构，从而使其能够突破局部最优，探索更广阔的优化地形。在内核运行时优化基准上的实验验证了这些增益：在 100 步的预算下，AdaExplore 在 KernelBench Level-2 和 Level-3 上分别实现了 3.12 倍和 1.72 倍的加速，且随着计算量的增加，性能持续提升。我们的代码已开源，访问 https://github.com/StigLidu/AdaExplore 获取。

## 1 Introduction
大语言模型（LLMs）已迅速演进为强大的编码智能体，在错误修复、代码重构和单元测试等任务上取得了优异表现（Chen et al.,2021 (https://arxiv.org/html/2604.16625#bib.bib6); Li et al.,2022 (https://arxiv.org/html/2604.16625#bib.bib2); Nijkamp et al.,2022 (https://arxiv.org/html/2604.16625#bib.bib28)）。近期研究进一步将 LLMs 扩展为工具增强型智能体，能够迭代式地合成并调试程序（Wang et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib34); Qian et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib32); Zhang et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib33)）。在本研究中，我们聚焦于 Triton 等底层编程框架中的 GPU 内核代码运行时优化（Ouyang et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib12); Li et al.,2025a (https://arxiv.org/html/2604.16625#bib.bib11)）。与仅关注功能正确性的传统代码生成不同，代码优化需在确保正确性作为硬性约束的前提下，对运行时性能进行优化。这种明确的性能目标为持续改进提供了天然的反馈信号。然而，尽管存在这一有利信号，搜索空间依然极具挑战。首先，如图 1 (https://arxiv.org/html/2604.16625#S1.F1)a 所示，可行性边界非常陡峭：语法、内存访问或并行化方面的微小错误往往会导致编译失败或运行时错误。由于底层编程语言训练数据的匮乏，这一挑战被进一步加剧，导致模型对有效实现的先验认知较弱，进而产生高比例的无效代码。其次，性能地形高度非线性和组合化（图 1 (https://arxiv.org/html/2604.16625#S1.F1)b），有意义的性能提升通常需要协同的结构变更，而非简单的本地编辑。此类变更通常涉及一系列相互依赖的修改序列，这使得内核运行时优化本质上成为一个长视野搜索问题。在实际操作中，专家级程序员通过迭代探索替代设计方案（如分块策略、内存布局和并行方案）来应对这一挑战，通常需要经过多轮试错与精细化调整才能实现高性能实现（Lim et al.,2017 (https://arxiv.org/html/2604.16625#bib.bib46)）。

**图 1 说明**：内核优化瓶颈示意图。(a) 由于训练覆盖范围有限，大多数生成的内核无效；(b) 内核细化可能陷入局部最优；(c) 我们的智能体 AdaExplore 能够从失败中学习技能以规避陷阱，并应用多样性保持搜索以探索全局最优。

在本工作中，我们将内核运行时优化视为在正确性约束下的搜索问题，其性能地形高度非线性和组合化。我们不依赖外部数据或微调，而是研究编码智能体如何通过累积的执行反馈和结构化探索来实现渐进式改进。这些观察结果凸显了两个核心挑战：(C1) 可行性探索：如何确保探索过程保持在由正确性约束定义的狭窄可行集内；(C2) 优化效率：如何在崎岖的优化地形中平衡全局探索与局部细化（Sutton et al.,1998 (https://arxiv.org/html/2604.16625#bib.bib14)）。这些挑战自然引出了两种分解思路：(i) 学习可复用的约束以维持在可行集内；(ii) 保持候选方案的多样性以探索崎岖的地形。为此，我们提出了 AdaExplore，这是一种基于 LLM 的内核运行时优化框架，其构建于两个互补机制之上：Adaptation（自适应）与 Exploration（探索）。在自适应阶段，智能体合成训练任务，并利用执行失败构建包含有效性规则的跨任务记忆库，以此引导后续代码生成走向可行集。经验表明，这显著提升了未见问题实例的正确率，且能泛化至不同的语言模型。在探索阶段，智能体在候选内核树上进行结构化搜索，维护多个候选项，并将多样化的解决方案轨迹作为上下文信号加以探索。它在小幅度的本地编辑与较大规模的结构变更之间交替进行，同时复用历史优秀候选项，从而实现突破局部最优的有效探索。与常见的测试时优化策略（包括迭代细化、并行采样和 OpenEvolve (Sharma,2025 (https://arxiv.org/html/2604.16625#bib.bib42))）相比，该方法在推理时扩展性上表现更佳，且随着搜索预算的增加，收益进一步提升（图 3 (https://arxiv.org/html/2604.16625#S4.F3)）。综合来看，这些机制帮助 AdaExplore 更可靠地定位有效内核，并在优化地形中更高效地搜索高性能方案。

我们的贡献总结如下：
- **基于失败的记忆机制**：提取可复用约束，在无需模型训练的情况下提升低资源场景下代码生成的有效性。
- **多样性保持的结构化搜索设计**：平衡本地细化与结构探索，实现高效的测试时扩展。
- 我们证明了结合这两种机制可获得最佳的整体性能。在 KernelBench 上，以 GPT-5-mini 为基础模型，在 100 步预算限制下，AdaExplore 在 Level-2 和 Level-3 上分别达到 3.12 倍和 1.72 倍的加速比。

**图 2 说明**：AdaExplore 内核运行时优化总览。该方法包含两个阶段：Adapt：将合成任务上的失败转化为跨任务记忆库，辅助生成正确内核。Explore：将候选内核组织为树状结构，在本地细化与结构重构之间交替，以搜索性能更高的方案。

## 2 Related Work
**LLM 代码生成。** Codex (OpenAI,2025 (https://arxiv.org/html/2604.16625#bib.bib1)), AlphaCode (Li et al.,2022 (https://arxiv.org/html/2604.16625#bib.bib2)), CodeLlama (Rozière et al.,2023 (https://arxiv.org/html/2604.16625#bib.bib3)), Seed-coder (Seede et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib4)), 和 Qwen-Coder (Hui et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib5)) 等通用代码模型在 HumanEval (Chen et al.,2021 (https://arxiv.org/html/2604.16625#bib.bib6)) 和 SWE-bench (Jimenez et al.,2023 (https://arxiv.org/html/2604.16625#bib.bib7)) 等基准上取得了优异成果，尤其是在训练数据充足的代表性语言上。在多语言代码生成方面，Multi-SWE-bench (Zan et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib8)) 和 HumanEval-XL (Peng et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib9)) 等基准突显了低资源编程语言日益受到关注的研究方向，并表明当前 LLM 在低资源或域特定语言上仍面临困难。

**自改进代码智能体。** 近期的自改进代码智能体主要可分为两类：跨任务适配与任务内优化。跨任务方法从过往试验中累积可复用的反思、技能或提示词（Shin et al.,2023 (https://arxiv.org/html/2604.16625#bib.bib30); Zhao et al.,2023 (https://arxiv.org/html/2604.16625#bib.bib31); Yuksekgönül et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib24); Agrawal et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib23); Zhang et al.,2025b (https://arxiv.org/html/2604.16625#bib.bib25); Wang et al.,2023 (https://arxiv.org/html/2604.16625#bib.bib58)）。此处，“任务”指同一领域内的独立问题实例，而非跨领域迁移。任务内方法利用环境反馈对单个问题的候选程序进行迭代细化或搜索，其思想可追溯至遗传编程（Koza,1992 (https://arxiv.org/html/2604.16625#bib.bib35)），并涵盖 LLM 指导的进化搜索方法如 FunSearch (Romera-Paredes et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib36)), AlphaEvolve (Novikov et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib37)), 和 CodeEvolve (Assumpção et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib39))，以及共进化与奖励驱动的方法如 CoCoEvo (Li et al.,2025b (https://arxiv.org/html/2604.16625#bib.bib40)。近期研究还将结构化搜索视为一种推理时扩展手段（Snel et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib56); Light et al.,2024 (https://arxiv.org/html/2604.16625#bib.bib57)）。我们的方法融合了上述两种视角：在搜索前跨任务提炼可迁移的失败模式，随后在每个任务内部采用多样性保持优化。

**GPU 内核生成的 LLM 智能体。** 我们聚焦于 Triton DSL，由 Tillet 等人 (2019 (https://arxiv.org/html/2604.16625#bib.bib10)) 引入，作为一种类似 Python 的语言用于开发 GPU 内核，并可 JIT 编译为 NVIDIA GPU 的 PTX 指令。尽管其语法高级，但编写高效的 Triton 代码仍需具备 GPU 架构的专业知识，包括线程束（warps）、共享内存和内存合并等概念。TritonBench (Li et al.,2025a (https://arxiv.org/html/2604.16625#bib.bib11)), KernelBench (Ouyang et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib12)), 和 Flashinfer-Bench (Xing et al.,2026 (https://arxiv.org/html/2604.16625#bib.bib27)) 等基准表明，当前语言模型在处理真实的 GPU 内核任务时依然吃力。为填补这一空白，近期的专用内核方法同时探索了智能体搜索与模型训练。Astra (Wei et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib52)) 将优化分解为多个智能体角色，而 AccelOpt (Zhang et al.,2025a (https://arxiv.org/html/2604.16625#bib.bib53)) 则将束搜索与针对新兴 AI 加速器的优化记忆相结合。同期工作 KernelSkill (Sun et al.,2026 (https://arxiv.org/html/2604.16625#bib.bib55)) 通过专家分析构建可复用技能，而我们则直接从执行反馈中自动提炼失败模式。在训练层面，Kevin (Barioni et al.,2025 (https://arxiv.org/html/2604.16625#bib.bib49)) 应用了多轮 GRPO，CUDA-L1 (Li et al.,2026 (https://arxiv.org/html/2604.16625#bib.bib50)) 使用带有加速评分示例的对比强化学习，而 CUDA Agent (Dai et al.,2026 (https://arxiv.org/html/2604.16625#bib.bib48)) 则在合成数据上扩展了智能体 PPO。我们的工作与该方向最为接近，但独特之处在于明确结合了可复用知识蒸馏与多样性保持优化。

## 3 Method
### 3.1 Task Setup
我们将内核运行时优化形式化为一个程序重写与优化问题。**输入**为一个原子函数（如矩阵乘法）的高级实现（如 Python 代码）；**输出**则为用特定底层语言（如 CUDA/Triton）编写的内核代码，该代码需满足 (i) 保持功能正确性 以及 (ii) 在目标硬件上最大化运行时性能。这使得该任务天然具有挑战性：智能体必须在保持高级程序语义的同时，发掘出能够满足硬件约束并实现强劲性能的底层实现。

### 3.2 Method Overview
我们的框架由两个组件构成：Adapt 与 Explore。Adapt 通过学习可复用约束来直接解决可行集探索问题，从而将优化过程限定在可行集内。通过在合成任务上运行智能体并收集执行失败信息，我们构建了一个紧凑的跨任务技能记忆库，其中包含关于何种情况容易导致内核无效的简单规则。该跨任务技能记忆库减少了推理期间的语法与执行错误，从而提升了生成准确率，并隐性加快了搜索速度。Explore 通过平衡候选方案的多样性与性能，帮助更有效地在优化地形中搜索高性能内核。我们将候选内核保存在树状结构中而非单一链式中，使搜索能够同时保留多个有前景的方向。每次扩展都会在微小的本地细化与较大的结构变更之间交替进行，结合近期探索历程与先前发现的优质内核来导航优化地形。两个组件的详细算法见附录 D (https://arxiv.org/html/2604.16625#A4)。

### 3.3 Adapt: Learning Skills from Failures
经验表明，内核生成失败通常源于少量反复出现的语法错误与结构约束（例如不支持的 Triton 操作、constexpr 违规等）。我们不依赖额外的模型训练，而是通过自我探索来调整模型对这些约束的认知。如图 2 (https://arxiv.org/html/2604.16625#S1.F2)a 所示，我们合成参考程序作为训练任务，并要求智能体实现对应的内核代码。通过对所产生的执行反馈进行总结聚类，我们将反复出现的失败模式提炼为系统提示词，构建成跨任务技能记忆库。

##### Task Synthesis
我们的流水线始于由标准算子组成的高级参考实现（如 PyTorch 程序）。我们使用一组不同于测试集的少量任务示例作为种子，随后复用并重组语言官方文档中描述的算子，以持续合成多样化的训练任务，而非依赖固定的人工精选集合。对于每个合成的任务，智能体会生成一个对应……