AsmEvo: Agentic Assembly-Level Optimization of AMD GPU Kernels with Functional Equivalence Verification —— 精读

论文链接:https://arxiv.org/abs/2608.20711

发表时间:2026 年 8 月(arXiv: 2608.20711v1,2026-08-21 提交,cs.CL)

发表机构:AMD(Advanced Micro Devices, Inc.)× 南方科技大学(深圳)。第一作者 Ji Liu(刘骥)领衔,署名作者共 21 人——20 人来自 AMD,Shengcai Liu 来自南方科技大学,通讯性质作者 Emad Barsoum 为 AMD Corporate VP/CTO(AI)。部分工作由 AMD 实习生在实习期间完成(论文脚注明确标注),是一次典型的「企业 + 高校」合作。

领域分类:cs.AI / cs.PL 交叉——GPU 内核优化、汇编级智能体搜索、AI 基础设施

一句话总结:不给源码、不给参考实现,只给一份已编译的 AMDGPU 二进制,智能体能不能把它改得更快且不出错?AsmEvo 给出了肯定答案——它把原始二进制当作唯一的「行为判据」,在恢复出的汇编上做 profiling 引导的局部编辑,只有差分验证通过的候选才被接受。


一、论文背景

1.1 GPU 内核优化的层级链条

一个机器学习模型在 GPU 上跑起来,中间要经历一长串「降级」过程:PyTorch 算子或 Triton DSL 代码,经过编译器(TVM、Triton 编译器)的调度选择、分块策略、内存排布,一层层 lowering 到 GPU 汇编(AMD 平台上是 AMDGCN 指令集),最终打包成可加载的 code object(HSACO/ELF 格式)。传统优化思路集中在链条上游:改源码、调调度、跑 autotuner——TVM 的 AutoTVM、Ansor、MetaSchedule 都是这个路线的代表。

但论文开篇点出一个被普遍忽视的事实:内核的最终性能不仅由高层 schedule 决定,还由 lowering 之后的低层决策决定——wait-counter(等待计数器)的放置、指令选择、寄存器分配、内核描述符、ABI 可见的元数据。这些决策发生在编译器后端,一旦编译完成,就「冻结」在二进制里,源码层面的工具再也摸不到。

1.2 源码不可得:部署环境的常态

更现实的约束是:在很多生产部署中,源码根本拿不到。论文列举了几类典型场景:

  • 推理服务栈只暴露一个编译好的 HSACO 文件;
  • JIT 编译缓存产物(比如 Triton 运行时生成的内核)只有缓存工件可查;
  • 供应商直接分发二进制库——AMD 的 AITer、rocBLAS、Composable Kernel,用户拿到的就是 code object;
  • vLLM、SGLang 这类推理框架里内嵌的 Triton 生成内核,部署时同样以编译产物形式存在。

这有点像买了一块封在表壳里的机械表:你无法拆开重新设计齿轮系(改源码),但表壳上仍露出几颗可以微调的游丝螺丝(汇编指令)。问题是——怎么在不打开表壳的情况下安全地调这些螺丝?

1.3 现有 LLM 内核优化器的盲区

近年来 LLM 内核优化很热:KernelBench 建立了「执行反馈迭代精炼」的范式,后续工作有强化学习路线(Kevin、CUDA-L1、Dr.Kernel、CUDA-Agent)、智能体进化搜索路线(Meta 的 KernelEvolve)、专用内核生成模型(KernelLLM)等。但论文指出它们共享两个隐含假设:存在可编辑的源码工件,且存在独立的参考实现用来验证正确性。

已编译的 code object 恰恰两条都不满足——原始二进制既是部署工件,又是唯一的行为基准。这就是 AsmEvo 要攻克的更严格设定。


二、论文定位与关联工作

2.1 研究版图上的四条线

把相关工作按「优化发生在哪一层、拿什么当正确性判据」整理成一张表,AsmEvo 的位置就清晰了:

技术路线代表工作操作层级正确性判据与 AsmEvo 的关系
LLM 源码级内核优化KernelBench、GEAK、Kevin、KernelEvolveCUDA/Triton/HIP 源码独立参考实现(如 PyTorch 算子)假设有源码与参考,AsmEvo 面向两者皆无的场景
张量编译器 / autotunerTVM、AutoTVM、Ansor、Triton调度 / IR框架级参考编译前或编译中起作用,编译后无能为力
汇编超优化STOKE、Souper、AlphaDev、CuAsmRLCPU 汇编 / LLVM IR / NVIDIA SASS指令规格或受限变换CuAsmRL 只做 NVIDIA SASS 指令重排;AsmEvo 面向 AMDGPU 且变换空间更宽
GPU 二进制工具NVBit、LuthierNVIDIA SASS / ROCm code object无(面向插桩)二者接受开销做 instrumentation;AsmEvo 追求更快且经验证的替换件

2.2 AsmEvo 的独特定位

论文自称(且据作者所知属实):AsmEvo 是第一个同时结合智能体搜索、ABI 与元数据保持的 AMDGPU code object 重建、以及对原始部署二进制差分验证的系统。它不替代编译器、autotuner 或推理库,而是补上「编译之后、部署之前」这最后一环。

值得注意的对比是 CuAsmRL(CGO'25):它把 NVIDIA SASS 指令调度建模成「汇编博弈」用强化学习求解,但变换被限制在指令重排,且依赖模型自身的正确性推理。AsmEvo 的候选编辑空间更大(指令替换、删除、插入、hoisting、交织等),正确性则完全交给外部确定性验证,模型只负责「提议」不负责「判定」。


三、问题定义

3.1 抽象:黑盒工件的性能优化问题

把具体场景剥掉细节,AsmEvo 处理的是一个相当一般的抽象问题:给定一个黑盒工件 $K_0$(已编译 AMDGPU code object),寻找优化版 $K'$,需同时满足三个约束:

  • R1(ABI 保持):$K'$ 保持符号、kernarg 布局、launch 语义与所有外部可见元数据——原始的 host 端 launcher 不改一行代码就能照常调用它。换句话说,优化件必须是「即插即用」的。
  • R2(功能等价):对所有被评估的输入 $x$,$K'$ 的输出必须与原始二进制的输出一致——余弦相似度 $c_x \geq 0.9999$、逐元素无穷范数偏差 $\leq 10^{-3}$(浮点输出),整数缓冲与不透明状态要求逐位一致,且越界保护 guard 必须完好。
  • R3(延迟改善):$K'$ 实测延迟更低。加速比定义为 $\hat{s}(K') = T(K_0)/T(K')$。

3.2 问题的难点所在

这个定义看似简单,实现起来有三重困难:

  1. 没有可编辑的表示。反汇编工具吐出的裸汇编缺少声明层(节区、符号、内核描述符、AMDGPU 元数据),直接改了也重汇编不回去。
  2. 没有验证用的输入。code object 既不暴露高层函数签名,也没有应用逻辑来构造指针密集的参数。怎么知道「喂给内核什么输入、怎么比对输出」本身就是难题。
  3. 错误信号极难察觉。一处局部汇编编辑可能破坏内存、违反 kernarg ABI、让元数据失同步,甚至「看似更快」仅仅是因为语义变了(少算了一半的数据当然快)。论文用了一个精辟的说法:优化信号必须不可作弊——快但错的候选,必须在计时之前就被拒掉。

一个类比:这相当于拿到了一份只有答案没有题目的考卷(二进制行为就是标准答案),要重写一份「过程不同但答案一致、且写得更快」的答卷。差分验证就是铁律——改前改后跑完全相同的输入,输出必须一致,浮点允许极小容差,整数与状态必须逐位相同。


四、问题解法

AsmEvo 是一条五组件流水线,外加一个长程智能体搜索驱动器。核心设计哲学是**「确定性控制器拥有所有正确性关键操作,LLM 只负责提议」**。

4.1 组件一:Code Object 恢复与往返保真

AsmEvo 从 AMDGPU ELF 容器中重建可编辑汇编 $s(K_0)$:恢复节区、符号、notes、内核描述符、AMDGPU 元数据和 AMDGCN 指令体,重建 .amdhsa_kernel 声明,并把 PC 相对寻址的控制流重新符号化——这样后续插入、删除、重排指令后分支目标仍然有效。

恢复之后有一道一次性往返保真门:把 $s(K_0)$ 重新汇编、重新链接,再与 $K_0$ 逐字节比对(仅屏蔽由链接决定的字段)。通过这道门,才证明后续编辑作用在一个忠实的、可重建的表示上。这相当于动手调表之前,先确认「拆下来再装回去,走时不变」。

4.2 组件二:两级差分 oracle

如何获得验证所需的输入与基准输出?AsmEvo 设计了双轨机制:

  • Tier 1(合成推断):当内核的 launch 结构能从描述符和元数据推断出来时,AsmEvo 推断参数偏移、标量类型、指针角色、缓冲大小和 launch 维度,然后生成确定性输入——覆盖代表性形状、步长、dtype、边界值和随机种子。特别巧妙的是,分配内存时在边界外放置金丝雀值,等价性检查由此不仅能发现输出分歧,还能发现越界写。
  • Tier 2(真实调度捕获):生产内核常用混合 dtype、非连续张量、block table、嵌套指针、应用自定义 workspace——元数据远不足以重构这些。AsmEvo 拦截一次真实的应用调度,记录 kernarg 缓冲、grid/workgroup 维度、引用的显存区域和调度前内存状态。原始对象先执行一次产生基准后状态;候选内核以完全相同的 launch 参数与内存内容重放这段调度。捕获参数里若含绝对设备指针,相关内存区域会在原虚拟地址恢复,嵌套指针无需重构应用层数据结构即可保持有效。

4.3 组件三:元数据感知的 ABI 保持重建

编辑后的汇编如何变回可加载对象?AsmEvo 重新扫描汇编中的 VGPR/SGPR/AGPR、LDS、scratch 用量,在内核描述符和 AMDGPU 元数据中重算资源字段后重新链接;同时冻结 kernarg 布局、kernarg 段大小、符号接口和 launch 语义——优化对象因此保持与原 launcher 的完全兼容。

若常规重建失败(对象含不支持或未完全恢复的构造),还有一条保守退路:原地字节 patch。但它被严格限制在「资源中性且体积不增」的编辑上,保留原始描述符与元数据;patch 出的候选照样过静态一致性、差分等价和性能三道门——退路改变的只是重建机制,不是接受标准。

4.4 组件四:Profiling 与热窗定位

大内核动辄数千条指令,全部塞进优化上下文既低效又容易诱发无关改动。AsmEvo 用硬件计数器、指令采样和静态分析定位 stall 主导的指令窗口(热窗),把热窗连同周边数据依赖与相关元数据交给搜索驱动器;提议的局部编辑拼接回完整汇编后,再做全内核重建与验证。热窗定位收缩了搜索上下文,但不放松全内核的接受标准——局部优化,全局把关。

4.5 组件五:门控验证 harness

每个候选按固定顺序过五道门:G0 汇编有效 → G1 静态 ABI/资源一致 → G2 功能等价(每个候选重复验证三次)→ G3 计时(预热后取 R 次事件计时中位数)→ G4 提交门。提交门采用方差感知阈值:

$$m^\star = (1 + \max(\epsilon, k \cdot c_v)) \cdot \hat{s}(K^\star)$$

其中 $c_v$ 是计时变异系数、$\epsilon = 0.002$、$k = 0.85$,另设 $s_{floor}$ 拒绝平凡重写。这防止测量噪声推动搜索前进。正确性严格先于计时——这是「不可作弊」的制度保障。

4.6 长程智能体搜索

验证 harness 接受任意候选生成器,AsmEvo 用长程 LLM 驱动器实例化它(实验中用 Claude Opus 4.8)。确定性控制器掌管时钟预算、工作区状态、重建与验证调用、提交决策、候选谱系与终止条件;LLM 分析 profiling 证据、选择瓶颈、提议局部汇编变换、解读结构化失败、在无效尝试后重规划。模型可以引导探索,但无权宣布正确或宣布更快。

几个值得细品的设计:

  • 验证谱系:被接受的候选构成树状谱系而非单一破坏性链条。每个节点记录父节点、验证加速比、改动窗口、资源用量与优化理由;工作者可以从当前最优分支,也可以回退到更早的验证版本。
  • 多起点探索(team 模式):规划者把互相正交的优化方向——延迟隐藏、依赖链削减、寄存器压力控制、访存重组、指令简化——分配给并行工作者(默认 4 个 GPU 工作者);起点选自多样化的已验证谱系节点,而不是所有工作者改同一份候选。
  • 验证式组合:两个各自有效的编辑拼在一起未必仍正确或仍有收益。集成器只考虑共享基底、改动区域不相交或互补的候选,增量应用、逐步重建验证;组合只有在全面过门且优于双亲时才提交。
  • 记忆与反停滞监督:有界运行记忆存储成功洞见、门控失败签名、重复候选哈希。控制器检测到长期无改善、重复提议或基础设施故障时,会请求「重定向」——汇总已试方向、标记穷尽假设,驱动器随之更换瓶颈目标或回溯;持续停滞则升级为更大规模多起点探索。

类比来说,LLM 像一位经验丰富但偶尔自信过头的调表师傅,而验证 harness 是旁边那台寸步不让的原子钟——师傅可以大胆尝试任何手法,但只有走时与原子钟对齐且更快的表才能出厂。


五、评估指标与实验证据

5.1 实验设置

  • 硬件:CDNA3 架构 AMD Instinct GPU——KernelBench 跑在 MI308X,AITer 与 Triton JIT 跑在 MI300X。
  • 预算:每内核固定 12 小时(0.5 天)墙上时钟,4 GPU 工作者多起点搜索。
  • 基准来源:30 个 KernelBench 内核(Level 1 / Level 2 各 15 个)——注意这些 $K_0$ 不是弱基线:用 AMD GEAK 风格智能体加 GPT-5.0 精炼 HIP 版本五轮、比对应 PyTorch 算子更快之后才编译交付 AsmEvo;生产内核(4 个 AITer code object + 4 个 vLLM/SGLang 的 Triton JIT HSACO)则原样取自部署环境,每个都配有真实调度捕获。
  • 指标:验证加速比 $\hat{s}(K') = T(K_0)/T(K')$($T$ 取预热后中位延迟);未修改二进制恒为 1.00× 参考。

5.2 主结果:KernelBench

切分内核数提升几何均值最大尝试/内核提交/内核最佳案例
Level 11515(100%)1.56×3.88×119.37.2Depthwise Conv2D
Level 21514(93%)1.17×2.55×156.46.8Matmul+Swish+Sum+GroupNorm
合计3029(97%)1.35×3.88×137.87.0Depthwise Conv2D

几个关键观察:

  • 收益不局限于单一算子类型。L1 中位数加速 1.31×,六个内核超过 1.5×,覆盖深度卷积(3.88×)、MinGPT-NewGelu 激活(3.82×)、RMSNorm 归一化(2.63×)、上三角矩阵乘(2.28×)。作者借此论证:单一固定的窥孔规则不可能覆盖这个算子谱系——这正是需要智能体搜索的理由。
  • 融合内核收益更集中。L2 中位数仅 1.10×,只有 Matmul+Swish+Sum+GroupNorm(2.55×)超过 1.5×;一个 GEMM+GroupNorm+Swish 流水线停在 1.00×。解释:融合内核暴露的编译后局部优化空间更小——但 2.55× 的最佳案例说明个别融合工件仍藏有大机会。
  • 有用的汇编重写非常稀疏:所有评估候选中只有 5.1% 成为验证提交。平均每个内核尝试约 138 次、提交约 7 次。这组数字从侧面印证了验证门的严格——搜索大部分时间在失败,而失败是有价值的反馈信号。

5.3 主结果:生产负载(MI300X)

8 个生产内核逐个报告(因 launch 语义与 oracle 需求各异,不做单一聚合):

来源内核加速比
AITeraiter_fmoe_fp8_blockscale_subgu2561.310
AITeraiter_fmha_hd128_fp8_causal1.007
AITernew_aiter_fmha_hd192x128_bf16_causal1.025
AITernew_aiter_fmha_bwd_hd128_bf16_causal_a321.049
vLLMgemm_a8w8_blockscale1.234
SGLangsglang_fused_moe_lora_kernel1.025
vLLMfused_moe_kernel1.343
SGLangawq_gemm_kernel1.133

AITer 家族几何均值 1.09×(最大 1.31×),Triton JIT 家族几何均值 1.18×(最大 1.34×,vLLM fused MoE)。八个生产内核全部提升。案例研究揭示了收益的具体来源:围绕长延迟访存的调度调整、wait-counter 放置、cache 提示与 buffer-load 变体、削减地址生成与类型转换工作——都是局部而非结构性的编辑。1.025× 到 1.343× 的跨度也说明:不能假设一个句法上合理的汇编变换对所有 JIT 内核都有效,每个重建出的 HSACO 都必须实测。

5.4 端到端可用性验证

论文没有止步于重放 harness 内的数字,而是验证了「即插即用」属性:KernelBench 派生的 HIP 内核直接替换 MI308X 原对象、跑同一 benchmark harness;AITer 替换件在原始 launch 配置下通过其单元测试;vLLM/SGLang 则替换缓存的 Triton HSACO、禁用重编译、跑框架测试——测得的加速作为真实替换工件依然成立。

5.5 诚实的边界声明

论文明确承认:AsmEvo 提供的是经验性差分等价而非形式化验证,保证范围限于被评估的输入与 launch 配置;新形状需要新的推断输入或调度捕获;harness 架构无关,但当前恢复与重建实现针对 CDNA 类 AMDGPU;复杂应用状态可能必须依赖真实调度捕获;可达收益取决于编译后剩余优化空间。这种克制在同类论文中并不多见。


六、效果优势的根源解释

6.1 为什么编译器之后还有 1.35× 的空间?

这是全文最值得咀嚼的问题。编译器不是已经「优化过」了吗?论文的答案隐含在背景与案例研究中,可以归纳为三层:

  1. 编译器是通用 pass 的集合,不为特定硬件配置与调用模式定制。同一个内核,在不同 GPU(MI300X vs MI308X)、不同 launch 几何、不同真实输入分布下,最优的指令调度与 wait-counter 布局并不相同。编译器只能取一个「平均合理」的保守选择,而 AsmEvo 拿到了确切的部署上下文。
  2. 低层决策空间在编译时难以穷尽。wait-counter 放置、寄存器分配、指令选择之间相互耦合,编译器后端在这些维度上做的是启发式决策。热 profiling 能定位编译器启发式「赌错」的具体位置——比如某条长延迟访存后的等待计数器放得太远,制造了本可隐藏的 stall。
  3. 多轮验证式组合能累积微小的局部收益。单次编辑可能只快 0.3%,但谱系树上「每次提交都过方差感知门」的组合机制,让几十个不冲突的小改动滚雪球。每内核平均 7 次提交乘起来的复利,正是几何均值的来源。

6.2 为什么「快但错」没有污染结果?

因为接受标准与提议机制彻底解耦。LLM 提议的候选没有任何特权,一律过 G0–G4 门;等价性在计时之前检查且重复三次;整数与不透明状态逐位比对、金丝雀捕捉越界写、commit 门要求超越「当前已验证最优」而非原始基线。5.1% 的提交率说明绝大多数候选死在了门里——这个看似难看的数字恰恰是系统可信度的来源。

6.3 为什么生产负载增益小于 KernelBench?

AITer 与 Triton 生产内核本身出自调优过的工具链(vendor 库、Triton 编译器),剩余空间天然更小;而 KernelBench 基线虽经 GEAK 精炼,毕竟只迭代了五轮。这符合直觉:优化空间的分布取决于上游工件的成熟度——论文也强调这一对比不能归因于上游编译器本身(内核算子与 launch 配置不同),只能说明 AsmEvo 对两类工件都能恢复出有用机会。


七、必要知识反推

如果想让一个团队复现或超越 AsmEvo,需要哪些知识储备?可以分三层反推:

7.1 领域层(AMDGPU 生态的硬知识)

  • AMDGPU ISA 与 AMDGCN 汇编:指令语义、hazard 规则、wait-counter 机制——这是智能体编辑的「棋盘」;
  • AMDGPU ABI 与内核描述符:kernarg 布局、VGPR/SGPR/AGPR 资源声明、launch 语义——这是 R1 约束的物理载体;
  • code object 格式(ELF/HSACO):节区结构、notes、元数据 schema——这是恢复与重建的对象。

7.2 方法论层(可迁移的方法)

  • 差分测试:以原工件为 oracle 的输入构造、金丝雀越界检测、容差分级(浮点余弦 + 无穷范数、整数逐位);
  • 热路径分析:硬件计数器 + 指令采样 + 静态依赖分析定位热窗,把全局问题局部化;
  • 长程智能体系统设计:确定性控制器与 LLM 提议器的职责分离、门控验证、验证谱系、多起点并行、验证式组合、反停滞监督——这套「学习引导探索 + 外部认证结果」的模式是当前 agentic 系统的成熟范式;
  • 测量纪律:预热后中位数、方差感知提交阈值、共享硬件上的噪声抑制。

7.3 工程层(落地条件)

  • 硬件与软件栈:MI300X/MI308X(CDNA3)、ROCm 工具链;
  • 生产内核库的获取与验证通道:AITer code object 的选取、vLLM/SGLang 推理时 HSACO 的采集、真实调度拦截(kernargs、launch 几何、显存快照);
  • 可观的算力预算:每内核 12 小时 × 4 GPU 工作者 × 38 个内核,再加 LLM API 成本(Claude Opus 4.8 驱动全程搜索)。

八、通用性灵感

8.1 黑盒工件优化范式

AsmEvo 最大的抽象贡献是把问题定义为「以行为等价为唯一约束的黑盒工件优化」。这个范式不限于 GPU 汇编:任何「源码不可得或与最终产物距离过远」的编译产物——WASM 模块、JVM 字节码、固件镜像、甚至 SQL 执行计划——理论上都能套用同一套「恢复可编辑表示 → profiling 定位 → 局部编辑 → 差分验证门控 → ABI 保持重建」的骨架。前提是三件事:工件可恢复为可重建的中间表示、存在可复现的执行 oracle、外部接口(ABI)可冻结。

8.2 Profiling 引导的局部编辑优于全局重写

热窗机制的价值有两层:工程上收缩了 LLM 上下文、降低了无关改动的概率;方法学上,「局部编辑 + 全局验证」的组合比「整体重写」在验证通过率与收益可归因性上都更优。这对所有代码优化智能体都有启发:先定位,再动刀,动刀后全局把关。5.1% 的提交率也提示后来者——在汇编这个层级,探索的代价结构是「大量廉价失败 + 少量高价值成功」,搜索架构必须为高失败率而设计(失败签名入库、谱系回退、多起点)。

8.3 验证门控:让学习系统安全探索的制度设计

AsmEvo 的 LLM 不能自证正确,一切由确定性门控裁决——「模型提议、外部认证」的分工是本文对 agentic 系统设计最有通用价值的示范。在任何高风险自动化场景(代码优化、系统运维、自动交易)中,把「探索的自由」交给学习系统、把「接受的权力」交给确定性验证器,既能享受学习的灵活性,又能守住安全底线。配套的方差感知提交阈值(拒绝噪声驱动的假进步)与验证谱系(可回退、可组合、不破坏)是让这套制度长期跑得住的工程细节。

8.4 一个务实的判断

AsmEvo 不是要取代编译器或 autotuner——它在「编译已完成、源码不可得」的最后一公里上工作,且收益分布高度依赖上游工件的成熟度。更合理的想象是:未来的推理服务栈把它作为部署前的例行「二进制体检 + 精修」环节,与源码级优化器形成上下游互补。当大模型推理成本成为全局约束时,这种「每个生产内核再榨 9%–18%」的能力,经济价值会非常直接。


附录:关键数字速查

项目数值
KernelBench 提升数29/30(97%)
KernelBench 几何均值 / 最大1.35× / 3.88×
Level 1 几何均值(15/15 提升)1.56×
Level 2 几何均值(14/15 提升)1.17×
AITer 几何均值 / 最大(4/4 提升)1.09× / 1.31×
Triton JIT 几何均值 / 最大(4/4 提升)1.18× / 1.34×
候选验证提交率5.1%
每内核搜索预算12 小时 × 4 GPU 工作者
等价阈值余弦 ≥ 0.9999,$\|\Delta\|_\infty \leq 10^{-3}$,整数逐位一致
提交门参数$\epsilon = 0.002$,$k = 0.85$,$s_{floor} = 1.0$
驱动模型Claude Opus 4.8(provider 默认解码)
KernelBench 基线生成GEAK 风格智能体 + GPT-5.0 精炼 HIP 五轮