AsmEvo: 带功能等价验证的AMD GPU内核智能汇编级优化

# AsmEvo: 通过功能等价验证实现AMD GPU内核的智能体级汇编优化
来源: https://arxiv.org/html/2608.20711  
杨璞源 郑荣章 王帆 王经林  
致谢：工作在AMD实习期间完成。  
Muhammad A. Awad Mortis Huang Andy Chang Zekai Li Zeping Li Zihao An Yue Liu Yuchen Yang Jianghui Wang Chushi Chen Ziqiong Liu Fuwei Yang Dong Li Wen Heng Chung Shengcai Liu Emad Barsoum

###### 摘要  
高性能机器学习系统越来越依赖于GPU内核，而这些内核的可编辑源代码往往不可用、是自动生成的，或距离最终机器代码过于遥远，无法暴露剩余的优化机会。现有的基于大语言模型的内核优化器和自动调谐器主要针对CUDA、Triton、HIP或张量程序源代码进行操作，并针对参考实现进行验证。我们研究一个更严格的场景：优化已经编译好的AMDGPU代码对象，其中部署的二进制文件是唯一的行为预言机。我们提出AsmEvo，一个用于AMD GPU内核的智能体级汇编优化器。给定一个AMDGPU代码对象K₀，AsmEvo会重建一个可重新汇编的表示，由一个长期智能体提议底层编辑，重建一个保持ABI（应用二进制接口）不变的优化后对象，并且仅在与K₀在相同启动配置下进行差异验证后才接受候选方案。AsmEvo结合了代码对象恢复、元数据感知的重建、基于性能分析的热点窗口编辑、正确性优先的计时验证以及保守的原地补丁回退机制。我们在各种AMD GPU内核上对AsmEvo进行了广泛实验。在MI308X上，AsmEvo改进了30个选定的KernelBench内核中的29个，达到1.35倍的几何平均加速和3.88倍的最大加速。在MI300X生产工作负载上，它改进了所有评估的AITer二进制文件以及vLLM/SGLang的Triton汇编内核，分别达到1.09倍/1.31倍和1.18倍/1.34倍的几何平均/最大加速，同时保持了功能等价性。

## 引言
现代机器学习工作负载依赖于GPU内核，这些内核的最终性能不仅由高层调度决定，还由降低表示后做出的底层决策决定：等待计数器放置、指令选择、寄存器分配、内核描述符以及ABI可见的元数据（Tillet, Kung, and Cox 2019 (https://arxiv.org/html/2608.20711#bib.bib26); Chen et al. 2018a (https://arxiv.org/html/2608.20711#bib.bib6); Zheng et al. 2020 (https://arxiv.org/html/2608.20711#bib.bib32)）。编译器（Chen et al. 2018a (https://arxiv.org/html/2608.20711#bib.bib6)）、张量程序自动调谐器（Zheng et al. 2020 (https://arxiv.org/html/2608.20711#bib.bib32); Chen et al. 2018b (https://arxiv.org/html/2608.20711#bib.bib8)）、领域专用语言（如Triton）以及供应商库（如AITer（AMD ROCm 2025 (https://arxiv.org/html/2608.20711#bib.bib4)）、rocBLAS（AMD ROCm 2016 (https://arxiv.org/html/2608.20711#bib.bib1)）和Composable Kernel（AMD ROCm 2019 (https://arxiv.org/html/2608.20711#bib.bib2)）都非常有效，但它们也创建了实际的界限。一旦内核成为AMDGPU代码对象（AMD ROCm 2023 (https://arxiv.org/html/2608.20711#bib.bib3)），剩余的优化机会可能无法通过源程序暴露——在许多部署中，源代码根本不可用。

这种无源码的场景在已部署的机器学习系统中很常见：服务栈可能只暴露编译后的HSACO（AMD ROCm 2023 (https://arxiv.org/html/2608.20711#bib.bib3)）、JIT缓存构件（Tillet, Kung, and Cox 2019 (https://arxiv.org/html/2608.20711#bib.bib26)）或供应商二进制文件，例如嵌入在vLLM（Kwon et al. 2023 (https://arxiv.org/html/2608.20711#bib.bib12)）和SGLang（Zheng et al. 2024 (https://arxiv.org/html/2608.20711#bib.bib33)）中的AITer代码对象或Triton生成的内核。重新优化已部署的构件很有吸引力，因为它针对的是GPU上实际运行的精确代码，但它也很脆弱：一个局部的指令编辑可能破坏内存、违反kernarg ABI、使元数据不同步，或者仅仅通过改变语义来显得更快（Villa et al. 2019 (https://arxiv.org/html/2608.20711#bib.bib27); Raayai-Ardakani et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib21)）。

最近基于LLM的内核系统（Wang et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib28); Zhang et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib31); Dai et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib9); Chen et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib7)）从高层规范（包括KernelBench（Ouyang et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib20)））生成和精炼CUDA、Triton、HIP或张量程序内核。它们假设存在可编辑的源代码构件和独立的参考实现。编译后的代码对象内核两者都不提供：原始二进制文件既是部署的构件，也是唯一的行为预言机。

本文探讨是否可以直接在*汇编级*优化内核，仅使用原始编译二进制文件来定义可接受的行为。我们提出AsmEvo，一个用于AMDGPU代码对象和恢复的AMDGCN汇编的编译后优化器。AsmEvo重建一个可重新汇编的表示，将热点指令窗口暴露给搜索驱动器，在保持启动器ABI的同时重建编辑后的候选方案，并且仅在针对原始二进制文件在相同输入和启动配置下进行差异验证后才接受候选方案。因此，优化信号是无法“作弊”的：快速但不正确的候选方案在计时之前就会被拒绝。

核心挑战是从一个没有高层签名的二进制文件中获取一个忠实的预言机。AsmEvo使用两级差异预言机。对于可以从元数据推断启动结构的内核，它合成受保护的输入并直接比较输出。对于混合数据类型、带步长、块表或指针丰富的生产内核，它捕获一次真实调度——kernargs、启动几何形状和引用的设备内存——并使用对浮点输出的容差感知检查和对整数及不透明状态的字节精确检查来重放候选方案。

AsmEvo通过一个元数据感知的重建路线将经过验证的编辑转换为可加载的对象：它重新生成描述符和AMDGPU元数据，同时冻结kernarg布局，以便优化后的对象仍然是一个直接替换品。仅当重建失败且编辑在资源上中性且不增加大小时，才使用保守的原地字节补丁。

性能分析和静态分析定位热点指令窗口，这些窗口被编辑、拼接回完整的汇编中、重建并重新验证。在此环境之上，AsmEvo使用一个在外部验证下的长期智能体搜索驱动器。确定性门控负责构建验证、资源检查、等价性、计时、提交阈值和血缘管理；智能体系划特定于架构的编辑，将失败的尝试作为记忆，探索多样化的已验证起始点，组合兼容的改进，并在搜索停滞时重新引导。这种分离使得基于学习的探索可以在有用的地方提供帮助，而每个被接受的结果都经过外部认证。

我们在各种AMD GPU内核上对AsmEvo进行了广泛实验。在MI308X上，AsmEvo改进了30个选定的KernelBench内核中的29个，达到1.35倍的几何平均加速和3.88倍的最大加速。在MI300X生产工作负载上，它改进了所有评估的AITer二进制文件以及vLLM/SGLang的Triton汇编内核，分别达到1.09倍/1.31倍和1.18倍/1.34倍的几何平均/最大加速，同时保持了功能等价性。这些结果表明，AsmEvo补充了源级编译器、JIT系统、自动调谐器和生产推理库。我们的贡献如下：

- 我们将*无源码、ABI保持的GPU内核优化*公式化为一个针对AMDGPU代码对象的、正确性优先的编译后问题，其中原始二进制文件是差异正确性预言机，每个报告的加速都必须通过外部验证。
- 我们构建了一个*二进制恢复、重建和差异验证流水线*：可重新汇编AMDGCN汇编的往返恢复、元数据感知的ABI保持重建、保守补丁回退以及结合合成启动推理与真实调度重放的两级预言机。
- 我们实例化了一个*在外部验证下的长期智能体搜索驱动器*：智能体系划特定于架构的编辑，将失败的尝试作为记忆，探索多样化的已验证起始点，组合兼容的改进，并从停滞中恢复，而确定性门控负责正确性和计时。
在KernelBench、AITer和vLLM/SGLang Triton JIT内核上的评估显示出一致的编译后收益。

## 相关工作
##### LLM内核优化与张量编译器。  
基于LLM的内核优化沿着几个互补的方向发展。KernelBench（Ouyang et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib20)）建立了一个执行引导的环境，其中生成的内核使用正确性和性能反馈进行迭代精炼。后续工作探索了具有测量奖励的强化学习（Baronio et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib5); Li et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib15); Liu et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib16); Dai et al. 2026 (https://arxiv.org/html/2608.20711#bib.bib9)）、在候选实现上的智能体和进化搜索（KernelEvolve Team, Meta Platforms 2026 (https://arxiv.org/html/2608.20711#bib.bib11); Lange et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib13)）以及专门的内核生成模型（Meta AI 2025 (https://arxiv.org/html/2608.20711#bib.bib19)）。相关的基准测试（Li et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib14); Lange et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib13)）进一步表明，性能优化仍然困难，且不充分的验证可能奖励快速但语义不正确的内核。

与此同时，张量编译器和领域专用语言优化调度、分块、内存移动和代码生成，包括Halide（Ragan-Kelley et al. 2013 (https://arxiv.org/html/2608.20711#bib.bib22)）、TVM（Chen et al. 2018a (https://arxiv.org/html/2608.20711#bib.bib6)）、AutoTVM（Chen et al. 2018b (https://arxiv.org/html/2608.20711#bib.bib8)）、Ansor（Zheng et al. 2020 (https://arxiv.org/html/2608.20711#bib.bib32)）、MetaSchedule（Shao et al. 2022 (https://arxiv.org/html/2608.20711#bib.bib25)）和Triton（Tillet, Kung, and Cox 2019 (https://arxiv.org/html/2608.20711#bib.bib26)）。尽管它们的搜索机制不同，但这些方法在编译之前或期间运行，并假设存在可编辑的源代码或中间表示，以及高层正确性参考。AsmEvo针对的是互补的编译后场景：其输入是生成的AMDGPU代码对象，而原始二进制文件充当行为预言机。

##### 汇编优化与二进制重写。  
超级优化器在指令或中间表示级别搜索更快的程序。STOKE（Schkufza, Sharma, and Aiken 2013 (https://arxiv.org/html/2608.20711#bib.bib24)）对x86汇编应用随机搜索，Souper（Sasnauskas et al. 2017 (https://arxiv.org/html/2608.20711#bib.bib23)）发现缺失的LLVM优化，AlphaDev（Mankowitz et al. 2023 (https://arxiv.org/html/2608.20711#bib.bib18)）学习底层的排序程序。最近基于学习的系统进一步优化汇编代码：Wei et al. 2025a (https://arxiv.org/html/2608.20711#bib.bib29)将强化学习应用于CPU汇编，而CuAsmRL（He and Yoneki 2025 (https://arxiv.org/html/2608.20711#bib.bib10)）将NVIDIA SASS指令调度表述为一种汇编游戏，并学习重新排序指令以提高吞吐量。这些方法假设面向CPU的规范或在具有约束转换的NVIDIA SASS调度上运行。AsmEvo则针对无源码的AMDGPU代码对象，并联合解决汇编恢复、ABI和元数据保持重建、更广泛的智能体提议编辑以及针对原始二进制文件的差异验证。

GPU二进制框架主要针对插桩。NVBit（Villa et al. 2019 (https://arxiv.org/html/2608.20711#bib.bib27)）修改NVIDIA SASS，而Luthier（Raayai-Ardakani et al. 2025 (https://arxiv.org/html/2608.20711#bib.bib21)）对加载的ROCm代码对象进行插桩和分析。对于这些用途，添加的开销是可以接受的，而AsmEvo寻求的是更快、经过验证的直接替换对象，同时保持启动器语义、资源声明、描述符和元数据。

##### 优化代码的正确性。  
编译器和翻译验证系统检查优化是否保持语义；例如，Alive2（Lopes et al. 2021 (https://arxiv.org/html/2608.20711#bib.bib17)）验证LLVM转换。程序等价性推理即使对强大的语言模型来说也很困难（Wei et al. 2025b (https://arxiv.org/html/2608.20711#bib.bib30)），这促使了基于执行的验证而非模型判断。现有的GPU优化器通常与源级或框架参考进行比较。对于仅HSACO的构件，这些参考不可用，因此AsmEvo将原始二进制文件视为差异预言机，并在计时之前检查等价性。其保证是经验性的，并且仅限于评估的输入和启动配置。

##### AsmEvo的定位。  
据我们所知，AsmEvo是第一个结合智能体搜索、ABI和元数据保持的AMDGPU代码对象重建以及针对原始部署二进制文件进行差异验证的系统。现有的源级内核智能体并非旨在处理仅HSACO的构件，或在没有此恢复和验证框架的情况下证明直接替换等价性。因此，我们在相同的框架内比较搜索驱动器，保持恢复、重建、正确性验证和提交门控固定不变。AsmEvo补充而非替代源级编译器、自动调谐器和推理库。

## 方法
### 概述
AsmEvo在没有源代码或独立参考实现的情况下优化编译后的AMDGPU代码对象。它将原始对象K₀视为差异预言机，并仅在ABI保持重建、功能等价验证和性能测量之后才接受优化后的对象K'。我们将恢复的汇编记为s(K)，评估的输入记为X，验证的加速比记为\(\hat{s}(K) = T(K_0) / T(K)\)。一个确定性控制器拥有所有正确性关键操作，而一个模块化的长期LLM驱动器从结构化反馈中提议编辑并重新规划。这种分离创建了一个可验证的环境，使得快速但不正确的候选方案无法推进搜索。

参见图1：AsmEvo结合了忠实的代码对象恢复、元数据感知的ABI保持重建，以及针对原始K₀的正确性优先于性能的评估。一个模块化的LLM驱动器通过多起点搜索、验证组合和防停滞重定向来提议编辑，而只有经过验证且延迟改进的候选方案进入血缘链。

### 问题形式化
给定K₀，AsmEvo寻求满足三个要求的优化代码对象K'。
R1（ABI保持）。K'保持符号、kernarg布局、启动语义和外部可见的元数据，允许原始主机启动器无需更改即可调用它。
R2（功能等价）。对于每个评估的输入x∈X，K'与原始二进制预言机O(x)=run(K₀, x)的输出匹配：eq(K', x) ≡ cₓ ≥ θ ∧ ‖Δₓ‖∞ ≤ τ ∧ ι(K', x)，其中cₓ = cos(K'(x), O(x))，Δₓ = K'(x) - O(x)。