跳转至

AVO: Agentic Variation Operators for Autonomous Evolutionary Search

会议: NeurIPS 2026
arXiv: 2603.24517
OpenReview: https://openreview.net/forum?id=MFTuhhh6v1
代码: https://github.com/Dao-AILab/flash-attention
领域: LLM Agent
关键词: 进化算法、智能体变异算子、GPU算子优化、注意力机制、Blackwell架构

一句话总结

将大模型从被动的候选解生成器跃升为具备自主长程规划与调试能力的进化变异算子(AVO),在 NVIDIA Blackwell B200 上经过 7 天自主迭代进化出吞吐超越 cuDNN(最高 3.5%)与 FlashAttention-4(最高 10.5%)的极致注意力算子。

研究背景与动机

大语言模型与进化搜索(Evolutionary Search)的结合近年来在算法发现与数学优化领域取得了突破性进展,如 FunSearch 与 AlphaEvolve 等系统已成功证明大模型在代码级解空间搜索中的潜力。然而,现有工作普遍将大模型严格局限在预设流水线内的候选生成步骤(Generate),外层的父代采样(Sample)、适应度评估(Evaluate)和种群管理依然由硬编码的启发式算法严格控制。在这种被动单轮调用范式下,大模型既无法在提交候选解前主动查阅微架构技术手册,也无法根据编译报错与性能 Profile 诊断日志进行针对性修正,更无法自主根据试错历史推演并调整长程优化战略。

对于 Attention 算子这类在现代 GPU 上已经被人类专家极致榨干性能的底层系统关键路径而言,这种割裂的单轮生成显得尤为无力。在 NVIDIA Blackwell(B200)架构上,无论是闭源的 cuDNN 还是开源巅峰 FlashAttention-4,都凝聚了系统专家数月乃至数年的手工微架构调优,包含了复杂的 Warp 特化协同、双级 Q-Tile 深度流水线与精细的张量内存加速器(TMA)异步编排。想要在如此逼近硬件理论极限的基线之上进一步发掘性能余量,必须依赖长期、多步的交互式工程闭环:反复研读 PTX 指令集与硬件白皮书、依据 Profiler 定位流水线停顿瓶颈、编写微架构改动并排查计算数值精度误差。

本文的切入角度是打破传统进化框架对大模型的流程束缚,将具备环境交互、持久记忆与工具调用能力的通用深度编码智能体直接赋权为进化变异算子本身。核心 idea:提出 Agentic Variation Operators(AVO),将采样、生成与评估验证收敛进自主编码智能体的单次变异闭环中,赋权智能体自主检索硬件知识库、编写底层 PTX/CUDA 内联代码、动态分析 Profiling 瓶颈并自适应反思调试,实现硬件底层算子的多天全自动演化。

方法详解

整体框架

传统 LLM 进化框架通常遵循固定的两阶段结构 \(\text{Vary}(\mathcal{P}_t) = \text{Generate}(\text{Sample}(\mathcal{P}_t))\),而 AVO 将其彻底重构为单一的自主智能体闭环 \(\text{Vary}(\mathcal{P}_t) = \text{Agent}(\mathcal{P}_t, \mathcal{K}, \mathbf{f})\)。智能体以历史进化谱系 \(\mathcal{P}_t\)、领域特定知识库 \(\mathcal{K}\)(包含 CUDA 编程指南、PTX ISA 规范、Blackwell 硬件架构白皮书与 FlashAttention-4 源码)以及多维评测函数 \(\mathbf{f}\) 为输入,在单一变异步内完成从架构瓶颈洞察到代码生成、编译校验、数值一致性判定与吞吐测试的全流程。

整个系统自上而下包含变异策略推演、内联内核微调修改、双重指标评估以及自监督防停滞干预机制,各模块紧密协作,形成多天连续演化的自主推进闭环。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:种子算子 x0 + 谱系历史 Pt + 硬件知识库 K"] --> B["变异策略推演<br/>检索PTX手册与微架构瓶颈定位"]
    B --> C["内核代码迭代编辑<br/>CUDA/PTX混合级微架构重构"]
    C --> D["双重指标评估验证<br/>数值正确性核对 + TFLOPS吞吐打分"]
    D -->|未通过精度或无增益| E["内部诊断反思循环<br/>读取编译器报错与Profiler停顿"]
    E --> C
    D -->|精度合格且吞吐刷新最优| F["提交新版本算子 xt+1<br/>写入Git并更新谱系 Pt+1"]
    F -->|探索停滞或连续循环| G["自监督干预机制<br/>重审演化轨迹并重定向探索空间"]
    G --> B
    F -->|继续下一轮自主变异| B

关键设计

1. 变异算子智能体化:统一采样、生成与验证的自主闭环

传统进化算法中变异算子被机械拆分为由框架调度的采样和被动单轮生成的两步,大模型无法控制自己的探索边界。AVO 将变异算子重新形式化为一个由通用编程智能体承载的自主回路: $\(\text{Vary}(\mathcal{P}_t) = \text{Agent}(\mathcal{P}_t, \mathcal{K}, \mathbf{f})\)$ 在单次变异步内,智能体不仅可以自由横向对比进化谱系 \(\mathcal{P}_t\) 中过往版本的性能 Profile 差异,还能主动查询知识库 \(\mathcal{K}\) 中的硬件约束(如寄存器文件容量、共享内存分配粒度及张量核心并发规则)。智能体在决定实施某项代码优化前,可以先进行轻量验证;若编译失败或数值精度比对未通过,智能体会主动读取底层错误输出,在单次变异步内多次自主修改调试,只有当候选内核完全满足数值精度且刷新最优吞吐时才向谱系提交新版本 \(x_{t+1}\),彻底摆脱了传统框架因无效候选而频繁中断的缺陷。

2. 无分支累加器重缩放与轻量内存栅栏:消除 Warp 发散与同步停顿

在 Online Softmax 算法中,当遇到更大的局部最大值时,输出累加器矩阵 \(O\) 需要按缩放因子进行实时修正。在前期内核版本中,这一重缩放逻辑依赖分支判断(检查 Warp 内是否有任意线程需要更新最大值),这虽然在理论上避免了无谓的乘法运算,但却在每个 Key-block 循环迭代中强行引入了昂贵的 Warp 同步开销,且迫使内存写回必须采用完全阻塞式的强内存栅栏(Blocking Memory Fence)。AVO 在 v20 版本中做出了关键底层洞察:在非掩码(Unmasked)迭代中,舍弃条件分支,将重缩放重构为纯投机式的无分支逻辑: $\(O \leftarrow O \times (\text{rescale\_factor} \text{ if updated else } 1.0)\)$ 通过谓词选择(Predicated Select)直接将不需缩放的因子置为 1.0,用极低代价的浮点乘 1 计算替换了分支发散;由于所有线程在此路径下完全对齐、严格共进,智能体得以安全地将阻塞式内存栅栏降级为轻量级的非阻塞内存栅栏(Non-blocking Fence),仅保留写顺序约束而不再全局停顿 SM 执行,为非因果注意力带来了高达 8.1% 的平均吞吐提升。

3. MMA 与 Correction 流水线重叠:打破双级 Q-Tile 串行瓶颈

Blackwell 上的 FlashAttention-4 采用了针对双 Q-Tile 的双级流水线(Dual Q-Stage),每个 Tile 内部包含计算注意力加权和的 PV GEMM,以及后续由专门 Correction Warp 执行的输出归一化。在 v29 版本中,智能体 Profiling 发现这两个 Tile 在 MMA 与 Correction 交界处被强制串行:Correction Warp 必须同步等待两个 Tile 的 PV GEMM 均全部执行完毕,才串行启动归一化计算,导致 Correction Warp 在第二个 GEMM 期间长期处于饥饿闲置状态。AVO 在 v30 版本重构了阶段间的同步信号流:一旦第一个 Q-Tile 的 PV GEMM 结束,立即解除对 Correction Warp 的屏障阻断,使其对 Tile 1 展开归一化处理,与 Tile 2 的 PV GEMM 在硬件执行单元上形成完全重叠的异步流水线,将原本串行的等待延迟彻底从关键路径中抹去。

4. 跨特化 Warp 组的寄存器精细重平衡:规避局部溢出并消除执行瓶颈

Blackwell SM 为每个线程块提供了固定配额(2048 个寄存器)的物理寄存器堆。FlashAttention-4 默认分配为:8 个 Softmax Warp 获得 192 个寄存器,4 个 Correction Warp 获得 80 个,其余 4 个 Warp 获得 48 个。智能体在演化至 v32 时通过底层 Profiling 指标敏锐发现:由于前述流水线重叠使得 Correction Warp 跃升至关键执行路径,其仅有的 80 个寄存器空间不足,导致大量局部变量发生 Register Spill,被迫溢出至高延迟的 Local Memory;而 Softmax Warp 凭借向量化打包指令,其峰值寄存器占用极低,192 寄存器配额存在严重冗余。AVO 在 v33 中将配额精准重平衡为 184 / 88 / 56(从 Softmax 组抽调 8 个寄存器补给 Correction 组),在不引发 Softmax 端任何退化的同时,彻底消除了关键路径上的溢出停顿。

5. 自监督干预机制:打破长程搜索的局部循环与停滞

长程自主搜索极易陷入局部最优停滞或重复无效微调的“死循环”。AVO 引入了全局自监督监控机制,周期性审视演化轨迹与提交历史;一旦探测到连续多次变异尝试均未能带来有效增益或陷入相似代码改动循环,该机制便会主动介入,对过往探索谱系进行全局复盘,强制向智能体注入全新的微架构探索方向(例如从指令调度层切换至内存流水线编排层),从而保障系统在无需人工干预的情况下实现长达 7 天的稳定正向演化。

训练与演化策略

AVO 采用基于真实硬件反馈的单谱系演化(Single-Lineage Evolution)机制。评测函数 \(\mathbf{f}(x_i) = (f_1(x_i), \dots, f_n(x_i))\) 在 NVIDIA B200 GPU 上针对不同序列长度(4096 至 32768,BF16 精度,Head Dim 128)分别采集端到端 TFLOPS。只有当候选算子与 PyTorch/cuDNN 参考实现的数值误差严格在允许公差内,且在测试集上的几何平均吞吐量不低于当前最优版本时,才触发 Git Commit 并纳入正式谱系。

实验关键数据

主实验

在 NVIDIA B200 上测试前向预填充(Prefilling)吞吐量(BF16 精度,Head Dim 128,总 Token 量固定为 32768),对比 NVIDIA 官方 closed-source cuDNN 9.19.1 与开源最优 FlashAttention-4(commit 71bf77c):

任务 / 序列配置 指标 (TFLOPS) AVO (本文) cuDNN 9.19.1 FlashAttention-4 相对 cuDNN 提升 相对 FA4 提升
MHA Causal (seq=4k, bs=8) 吞吐量 1668 ~1612 ~1510 +3.5% +10.5%
MHA Causal (seq=8k, bs=4) 吞吐量 1655 ~1625 ~1545 +1.8% +7.1%
MHA Causal (seq=16k, bs=2) 吞吐量 1640 ~1620 ~1560 +1.2% +5.1%
MHA Causal (seq=32k, bs=1) 吞吐量 1628 ~1622 ~1550 +0.4% +5.0%
MHA Non-Causal (seq=16k, bs=2) 吞吐量 1645 ~1615 ~1620 +1.8% +1.5%
MHA Non-Causal (seq=32k, bs=1) 吞吐量 1638 ~1600 ~1610 +2.4% +1.7%
GQA Causal (32Q4KV, seq=32k) 吞吐量 1610 ~1505 ~1473 +7.0% +9.3%
GQA Non-Causal (32Q8KV, seq=32k) 吞吐量 1620 ~1528 ~1550 +6.0% +4.5%

消融实验

针对智能体自主探索出的代表性微架构优化,在前后相邻版本间展开几何平均吞吐(Geomean TFLOPS)单点消融:

优化项 版本演进 Non-Causal 提升 Causal 提升 核心机理说明
无分支累加器重缩放与轻量内存栅栏 v19 → v20 +8.1% +1.6% 消除 Warp 条件发散,非掩码路径换用非阻塞内存栅栏
Correction/MMA 流水线异步重叠 v29 → v30 +1.1% +0.4% Tile 1 归一化与 Tile 2 PV GEMM 并行,消除空等停顿
跨特化 Warp 组寄存器精细重平衡 v32 → v33 +2.1% ~0.0% 从 Softmax 抽调 8 寄存器补给 Correction,消除 Spill 溢出

关键发现

  • 离散跃迁胜过连续渐进:全流程探索了超过 500 个优化方向,共沉淀 40 个正式提交版本。算子性能提升并非均匀微调,而是呈现“跃迁-平台期”特征,其中 5 次重大突破(QK-PV 交错、单趟 Softmax、无分支重缩放、流水线重叠、寄存器重平衡)贡献了绝大部分增益。
  • 极致优化的零样本迁移能力:在 MHA 上进化出的微架构创新具备高度鲁棒性与泛化能力。将最终 MHA 算子迁移适配至 GQA(Qwen3-30B 与 Qwen3-8B 头结构)仅耗费智能体 30 分钟无人工干预的代码重构,在因果与非因果场景下分别实现了相对 cuDNN 7.0% 与 6.0% 的性能飞跃。
  • 软硬件协同理解的深层体现:消融数据显示,无分支重缩放使 Non-Causal 获得了 8.1% 的巨大加速,但在 Causal 场景下收益为 1.6%,原因在于 Causal 存在大量被遮蔽的 Key-block 必须走保守分支路径。智能体能够精准识别这一边界并在安全区域实施激进优化,展示了深度的微架构推理能力。

亮点与洞察

  • 从候选解生成器跃迁为自主变异算子:AVO 首次将 LLM 的角色从传统进化框架内部受限的单轮生成模块,跃升为端到端掌控采样、编码、剖析、验证的自主智能体回路。
  • 软硬件边界的多层次协同推演:智能体发掘的改进并非浅层的代码语法变换,而是跨越了 PTX 谓词指令选择、内存一致性模型(内存栅栏降级)、异步流水线重组以及硬件寄存器物理分配的多维联合优化。
  • 长程进化自监督防卡顿机制:通过外部自监督检查点动态识别搜索停滞并注入全局重定向提示,确保了全自动演化在 7 天跨度内持续产出高质量突破。

局限与展望

  • 单谱系演化尚未利用分支并行:目前实验仅在单一谱系(Single-lineage)下验证了算子性能,尚未与经典的 MAP-Elites 或分布式多岛种群结构结合,探索广度仍有扩展空间。
  • 硬件与算子单一性约束:全部搜索与评测仅集中在 NVIDIA Blackwell B200 架构的 Forward Attention 上,未来需要进一步向反向传播、长上下文解码以及 AMD/TPU 等非 CUDA 异构芯片泛化。
  • 计算资源开销与探索效率:7 天持续占用顶级 B200 硬件资源并探索了 500 多个内部方向,若能引入基于历史经验的轻量级性能代理模型,有望显著压降编译与硬件实测的时间与能耗成本。

相关工作与启发

  • vs FunSearch / AlphaEvolve: FunSearch 与 AlphaEvolve 将 LLM 限制在框架预设的单轮生成环节,由外层启发式策略负责采样和评估;AVO 则将采样、工具查阅、代码修改与 Profiling 闭环统筹在自主智能体内部,能解决需要多步深度调试的微架构难题。
  • vs LoongFlow / TTT-Discover: LoongFlow 采用固定的 Plan-Execute-Summarize 机械流程,TTT-Discover 依靠测试时梯度更新生成策略;AVO 赋予智能体任意调用 Profiler、PTX 文档与历史谱系的完全能动性,支持在长程搜索中动态漂移策略。
  • vs FlashAttention-4: FA4 是顶尖专家数月手工打磨的极致算子;AVO 以 FA4 为基础底座,通过全自动多层微架构协同挖掘出了专家未曾触及的剩余硬件空间,展示了 AI 辅助系统优化的巨大潜力。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 颠覆了传统 LLM-in-the-loop 进化搜索范式,将 Agent 确立为自包含的进化变异算子。
  • 实验充分度: ⭐⭐⭐⭐⭐ 7天连续自主演化、500+方向尝试、40个迭代版本,在最高端B200芯片上正面击败cuDNN与FA4。
  • 写作质量: ⭐⭐⭐⭐⭐ 问题定义直击底层系统与微架构痛点,演化轨迹与消融机理剖析极其详实透彻。
  • 价值: ⭐⭐⭐⭐⭐ 为未来顶尖系统级基础软件与高性能算子的全自动自主演进开辟了极具说服力的全新技术路径。