Command Palette
Search for a command to run...
AVO:面向自主进化搜索的智能体变异算子
AVO:面向自主进化搜索的智能体变异算子
摘要
智能体变异算子(Agentic Variation Operators,AVO)是一类新的进化变异算子,用自主编码智能体取代经典进化搜索中固定的变异、交叉和手工设计的启发式方法。AVO 并不将语言模型局限于在预设流水线中生成候选解,而是将变异实例化为一个自我引导的智能体循环,该循环可以查阅当前谱系、领域特定知识库和执行反馈,以提出、修复、评判和验证实现编辑。我们在 NVIDIA Blackwell(B200)GPU 上,针对 AI 领域中被优化得最为激进的注意力核这一目标,对 AVO 进行了评估。在多头注意力上经过 7 天的持续自主进化后,AVO 发现的核在所有评估配置下,性能最高比 cuDNN 提升 3.5%,比 FlashAttention-4 提升 10.5%。所发现的优化可以轻松迁移到分组查询注意力,仅需 30 分钟的额外自主适应,便取得了比 cuDNN 最高提升 7.0%、比 FlashAttention-4 最高提升 9.3% 的增益。这些结果共同表明,智能体变异算子通过将智能体从候选解生成器提升为变异算子,超越了此前的 LLM 在环进化流水线,并能够发现对性能至关重要的微架构优化,从而在当今最先进的 GPU 硬件上产生超越顶尖专家工程实现的注意力核。
一句话总结
NVIDIA研究人员提出了Agentic Variation Operators (AVO),即自主编码agent,它们通过一个循环咨询谱系、知识库和执行反馈来提出、修复、评判和验证实现编辑,从而取代固定的进化算子,并在Blackwell B200 GPU上经过7天连续进化后,发现的多头注意力核函数性能比cuDNN最高提升3.5%,比FlashAttention-4最高提升10.5%,且快速迁移至分组查询注意力仅需30分钟适配,分别获得最高7.0%和9.3%的增益。
核心贡献
- Agentic Variation Operators (AVO) 用自主编码agent取代固定的进化变异启发式,将采样、生成和评估整合到一个自主引导的循环中。
- 在NVIDIA Blackwell GPU上应用于多头注意力时,AVO在七天的连续进化后发现核函数,性能比cuDNN最高提升3.5%,比FlashAttention-4最高提升10.5%。
- 所发现的优化在额外30分钟适配后迁移至分组查询注意力,相比cuDNN最高提升7.0%,相比FlashAttention-4最高提升9.3%。
引言
作者针对在进化代码优化中使用大语言模型的挑战,此前系统将LLM限制为固定流水线中的单轮候选生成器。这限制了进行迭代、多步骤工程的能力,而这是超越高度调优实现(如GPU上的注意力核函数)所必需的,后者需要研究硬件文档、性能分析、测试和修改。为克服此问题,他们引入了Agentic Variation Operators (AVO),它用自主引导的编码agent取代整个变异步骤,该agent自主查阅知识、评估候选方案并在长时间范围内优化其方法。在NVIDIA Blackwell B200 GPU上应用于多头注意力时,AVO发现了微架构优化,性能比cuDNN最高提升3.5%,比FlashAttention-4最高提升10.5%,并且所学的技术在仅30分钟的额外自主适配后即可迁移至分组查询注意力。
方法
作者首先详述了现代GPU上注意力计算的复杂性。给定查询、键和值矩阵 Q,K,V,注意力计算为 O=softmax(QK⊤/d)V。虽然FlashAttention通过顺序处理键块避免了实例化完整的得分矩阵,但NVIDIA Blackwell架构上的最新核函数(如FA4)采用了warp specialization。不同的warp组并发处理矩阵乘法、在线softmax、输出重新缩放和数据移动。优化此类高度调优的核函数需要深厚的硬件专业知识与大量性能分析。
为实现此优化,作者提出了Agentic Variation Operators (AVO)。传统进化搜索将变异算子分解为单独的采样和生成步骤,将LLM限制在固定流水线中。AVO用单次自主agent运行替代这一分解:
Vary(Pt)=Agent(Pt,K,f)其中 Pt 是解的完整谱系及其得分,K 是领域特定的知识库,f 是评分函数。
AVO系统的输入包含三个主要组成部分。种群 Pt 包含先前CUDA核函数实现的谱系及其得分。知识库 K 包含CUDA和PTX文档、Blackwell架构规范以及参考GPU核函数代码库。评分函数 f 根据数值正确性和TFLOPS吞吐量评估候选方案;任何未通过正确性检查的候选方案无论吞吐量如何均得零分。
该方法的核心是AVO主Agent循环,它从当前谱系 Pt 产生新解 xt+1。此自主循环涉及多个相互关联的阶段。在规划阶段,agent利用知识库和先前实现提出编辑建议。在实现阶段,它应用代码修改。评估阶段使用评分函数 f 测量性能。若出现问题,错误修复阶段允许agent诊断、修复并调整其计划。agent经常检查多个先前实现以识别瓶颈,并在实施优化之前查阅文档以理解硬件约束。此编辑-评估-诊断循环重复进行,直到提交满意的 xt+1。新版本仅在通过正确性检查且其基准得分匹配或优于最佳已提交版本时才会被持久化。
为确保长时间自主优化过程中的持续进展,作者引入了AVO Supervisor Agent。该监督者监控主agent是否停滞,例如耗尽某个探索方向或陷入无效的编辑循环。一旦触发,监督者会审查整体进化轨迹,并提供条件干预以引导搜索转向新的候选优化方向。
该系统以持续进化循环的方式运行。每个已提交版本均作为git提交持久化并附带其得分,从而在整个过程中保持完整的状态连续性。在他们的评估中,一次为期7天的运行产生了最终的多头注意力核函数,跨越40个连续版本,主agent自主决定何时尝试新的优化或切换策略,而监督者在停滞期间进行干预。
实验
评估使用agentic variation operator (AVO) 在NVIDIA B200 GPU上进化CUDA注意力核函数,并在多头和分组查询注意力配置下与cuDNN和FlashAttention-4进行基准对比。AVO自主发现了一些优化,这些优化相对于两个基线均带来一致的吞吐量提升,其中最大的改进来自无分支重新缩放、流水线重叠和寄存器重新平衡等架构变更,而后续版本则提供细粒度调优。agent在不到一小时内成功将其进化的MHA核函数适配到GQA,表明所发现的优化可泛化到原始搜索空间之外。总体而言,实验表明,agent驱动的进化能够通过迭代性能分析和代码修改,联合推理多个硬件子系统,从而产生专家级的核函数优化。
agent发现的三个核函数优化各自提升了吞吐量,增益因注意力类型而异。无分支累加器重新缩放在非因果注意力上带来了8.1%的大幅提升,但在因果注意力上仅提升1.6%,因为无分支路径仅适用于未掩码的键块。流水线重叠和寄存器重新平衡带来了进一步的增益,主要针对非因果场景,其中校正warp位于关键路径上。无分支累加器重新缩放在非因果注意力上提供了8.1%的几何平均吞吐量增益,但在因果注意力上仅1.6%,因为该优化仅限于完全未掩码的键块迭代。跨warp组的寄存器重新平衡在非因果注意力上带来2.1%的提升,对因果注意力影响可忽略,因为校正warp仅在非因果流水线中与第二个GEMM并发运行。
agent发现的核函数优化提升了吞吐量,增益集中在非因果注意力上。无分支累加器重新缩放对非因果场景非常有效,但对因果注意力益处有限,因为它仅适用于未掩码的键块。额外的流水线重叠和寄存器重新平衡进一步提升了非因果性能,其中校正warp位于关键路径上,而因果注意力则改善甚微。