斯坦福与亚马逊联手打造AI芯片"自动调音师"

这项由斯坦福大学计算机科学系与亚马逊云服务（AWS）联合开展的研究，发表于2026年第九届MLSys机器学习与系统国际会议（MLSys 2026），会议地点位于美国华盛顿州贝尔维尤。论文编号为arXiv:2511.15915v2，发布于2026年4月15日，感兴趣的读者可通过该编号检索完整论文。

一、为什么一块芯片发挥不出应有的力量

每当我们使用ChatGPT、Gemini或任何大型AI工具时，背后都有成千上万块专用芯片在疯狂运转。这些芯片被称为AI加速器，它们就像一台台精密的乐器，天生具备演奏高难度乐章的能力。然而，乐器再好，如果演奏者不懂得如何驾驭它，也只能弹出平庸的音符。

在AI芯片的世界里，这个"演奏者"就是所谓的**内核程序**（kernel）——一段专门告诉芯片如何处理矩阵乘法、注意力机制等计算任务的底层代码。内核写得好，芯片可以发挥出接近理论极限的性能；内核写得差，再贵的芯片也只能跑在30%甚至更低的效率上。问题在于，针对每一块新芯片写出高质量的内核，需要具备深厚专业知识的硬件工程师花费大量时间和精力去摸索调试。英伟达H100芯片于2022年发布，工程师们花了将近一年时间才把注意力机制的内核调到理论峰值的37%，又花了一年多才勉强接近85%。

亚马逊自研的Trainium芯片是这个故事的主角背景。它是一款专门用于训练大型AI模型的加速器，拥有独特的硬件架构，对应的编程接口叫做神经元内核接口（NKI，Neuron Kernel Interface）。由于这套接口相对较新，工程师们缺乏成熟的调优经验和现成的优化套路，很多内核程序都还运行在远低于硬件极限的水平上。

研究团队注意到这个困境，提出了一个大胆的想法：既然大型语言模型（LLM，就是ChatGPT那类会写文章、能聊天的AI）在理解和生成代码方面已经颇有造诣，能不能让AI来帮AI芯片调音？更进一步，能不能让这个AI调音师在工作中不断积累经验，越调越好，最终不需要任何人工干预就能自主掌握优化技巧？

这就是AccelOpt（Accelerator Optimizer，加速器优化器）诞生的初衷。

二、AccelOpt是什么：一位会自我进化的调音师

以调钢琴为比喻来理解AccelOpt会非常直观。一位顶尖钢琴调音师第一次面对一台从未见过的新型钢琴，他会怎么做？他会先试弹几个音，听听哪里跑调了，再根据经验判断可能的原因，然后动手调整，反复验证。更重要的是，每次成功调好一台琴，他都会把这次的经验记录下来——"这个型号的琴，中音区的弦往往偏紧，要用某种特定的手法处理"——下次遇到类似的琴，他就能更快更准地找到问题所在。

AccelOpt做的事情与此如出一辙。它面对的"钢琴"是Trainium芯片上运行的NKI内核程序，面对的"调音任务"是让这些程序跑得更快、更接近芯片的理论极限。整个系统由三个相互协作的AI智能体组成，分别扮演不同的角色，再加上一套"束搜索"机制和一个持续更新的"优化记忆库"，共同构成了这台自动调音机器。

三个智能体各司其职，形成一条完整的调音流水线。**规划师**（Planner）负责分析当前内核的运行剖析数据——也就是芯片在执行这段代码时，哪个部分耗时最长、哪里在白白浪费资源——然后提出一个具体的优化方向。**执行师**（Executor）拿到规划师的方案后，动手将其转化为实际的代码改动，真正生成一个新版本的内核程序。**总结师**（Summarizer）则在一次成功的优化完成之后，把这次的经验提炼成可以复用的通用规律，存入记忆库，供未来的任务参考。

三、束搜索：不把宝押在一次尝试上

光有三个智能体还不够。一个经验丰富的调音师知道，面对一台复杂的钢琴，光靠一次尝试找到最优解几乎是不可能的。他可能需要尝试多种不同的调整思路，分头推进，看看哪条路走得最顺。

AccelOpt采用的**束搜索**（beam search）机制正是这个道理。在每一轮迭代中，系统不是只保留一个最新的内核版本，而是同时维护着一批表现最好的候选内核——就像调音师同时在脑子里保留着几套并行的方案。具体来说，系统会给每个候选内核生成若干条不同的优化方向（由规划师提出），再对每个方向尝试若干次代码实现（由执行师完成），最后从所有这些新生成的版本中，挑选出表现最佳的一批，作为下一轮迭代的起点。

这套机制的好处在于，它避免了"走进死胡同"的风险。如果某次优化尝试走错了方向，损失的只是其中一个分支，其余分支还在继续前进。随着迭代轮次的增加，整个候选集的整体质量会稳步攀升。研究团队的实验也验证了这一点：束搜索比单纯地重复采样同一个提示词效果要好得多，因为它真正实现了从优秀走向更优秀，而不是在原地反复打转。

四、优化记忆库：经验比黄金更值钱

束搜索解决了"怎么探索"的问题，但还有另一个挑战：AI系统在这次优化任务中学到的东西，能不能在下一次任务中发挥作用？

这就是**优化记忆库**（optimization memory）的用武之地。它本质上是一个动态更新的经验档案，记录着每次成功或失败的优化尝试中，哪段代码的哪种改动带来了显著的速度提升或下降。更重要的是，总结师会把这些具体的代码改动抽象成更通用的规律，比如"把不变的矩阵转置操作移到循环外面，可以避免重复计算"。

记忆库的设计有几个值得关注的细节。一方面，它不仅记录"成功案例"（慢变快），也记录"失败案例"（快变慢），因为知道什么不该做有时候和知道什么该做同样重要。另一方面，为了防止记忆库被重复的、相似的经验刷屏，系统在选取经验条目时会刻意保持多样性——不同的候选内核、不同的优化方向，各自贡献一个最有代表性的经验点。

记忆库有两个关键参数控制其行为：ExpN控制库的容量（能保存多少条历史经验），TopK控制每轮迭代能往库里新增多少条新经验。研究团队发现，增大容量（ExpN）比增大每轮更新量（TopK）更划算——保留更多历史沉淀的价值，高于在每轮迭代中急着塞入更多新经验。这就好像一位调音师，积累十年工龄的深厚阅历，比刚入行时每天快速记笔记要更有价值。

五、NKIBench：给调音师设计一套考题

研究团队在评估AccelOpt时遇到了另一个问题：没有合适的标准考题。已有的AI芯片内核基准测试，要么不包含NKI程序，要么只是简单地比较"优化后比优化前快了多少倍"，却无法告诉你这个内核距离芯片的极限还有多远的路要走。

于是研究团队自己动手构建了**NKIBench**，这是第一个专门针对Trainium芯片NKI内核的基准测试套件。所有14个测试任务都来自真实的大模型工作负载，包括DeepSeek、Qwen3、Falcon等主流模型中涉及的矩阵乘法、批量矩阵乘、注意力机制、RMSNorm归一化、LoRA微调、Mamba序列模型等关键算子。

NKIBench的一个重要创新在于，它为每个任务计算了硬件理论峰值性能，并用实际运行时间占峰值性能的百分比来衡量内核的优化程度。这就好比一辆赛车在跑圈，不是只看它比上一版车快了几秒，而是直接告诉你它已经跑到了这条赛道理论最快圈速的多少百分比。这个指标更能客观反映优化工作还有多大的提升空间。

理论峰值的计算基于芯片的三个核心瓶颈：HBM带宽（芯片从外部内存读写数据的速度）、张量引擎算力（专门做矩阵乘法的计算单元的速度）、向量引擎算力（做其他数学运算的计算单元的速度）。具体公式是取这三个瓶颈中最紧的那个作为理论最快速度，然后用实际延迟与之相比。

六、调音的成果：从49%到61%

经过系统性的测试，AccelOpt的表现相当令人满意。在Trainium 1芯片上，NKIBench测试集的平均峰值吞吐量占比从优化前的49%提升到了61%；在更新的Trainium 2芯片上，则从45%提升到59%。换句话说，这批内核程序在优化后，能榨取出芯片更多的潜力，平均提升幅度在12到14个百分点之间。

更有趣的是与顶尖商业模型的对比。Anthropic公司的Claude Sonnet 4是当前公认的最强代码生成模型之一，研究团队也用它来做同样的内核优化任务（采用重复采样的方式，不断生成新版本直到预算用完）。最终，AccelOpt在使用开源模型（Qwen3-Coder-480B作为执行师，gpt-oss-120b负责其余角色）的情况下，达到了与Claude Sonnet 4相近的优化效果，但花费的API调用费用仅为后者的二十六分之一。

这个对比背后有一个很有意思的发现：Claude Sonnet 4的"重复采样"策略，本质上是在同一个起点反复尝试，就像一个调音师一遍遍地试弹同一个音，期望某次偶然能找到最优解。而AccelOpt的束搜索则是从优秀的版本出发继续精进，每一步都站在前一步的肩膀上。这种迭代式进化的策略，让开源模型在成本上大幅领先的同时，仍能保持竞争力。

七、AccelOpt发现了哪些具体的优化技巧

研究团队特别列举了AccelOpt自主发现的几类典型优化，展示了这套系统的实际能力深度。

第一类是**窥孔优化**（peephole optimization），也就是对局部代码片段进行精简和替换。比如，AccelOpt会自动识别出`θ_t–1 – γλθ_t–1`可以化简为`(1–γλ)θ_t–1`，从而把一次减法和一次乘法合并成一次乘法，减少计算量。它还能识别`reciprocal(sqrt(...))`这个组合可以被替换为更高效的`rsqrt(...)`单一指令，减少中间临时变量的创建。对于SiLU激活函数，AccelOpt发现`x/(1+e^{-x})`可以改写为`x·sigmoid(x)`，从而调用NKI专门优化过的sigmoid指令，获得更高效率。

第二类是**循环变换优化**，这类优化需要更深层次的推理。研究团队展示了一个BatchMatmul加Softmax融合算子的优化案例，非常能说明问题。初始版本的内核因为中间变量`v`和`p`需要跨越两个循环存活，被迫溢出到芯片外部的慢速内存（HBM）中，造成大量额外的数据搬运开销，延迟高达12毫秒，向量引擎利用率只有46%。AccelOpt在第一步发现了这个溢出问题，提出重新计算`v'`来消除溢出，延迟降到8.2毫秒，但代价是引入了额外的矩阵乘法计算。系统没有满足于这个结果，继续推进，在下一步彻底重构了循环结构，去掉了重新计算和多余的外层循环，最终版本延迟降到6.4毫秒，向量引擎利用率跃升至84%。

第三类是**循环不变量外提**（loop invariant code motion）。这个概念用大家熟悉的场景来说：假设你要做20道数学题，每道题开始都要先查一遍同一张乘法表——合理的做法显然是只查一次，把结果记在手边。AccelOpt在Transpose+Matmul的内核中发现，LHS矩阵的转置操作在每次外层循环迭代中都被重复执行了16次，而这个转置结果其实从不改变。系统自动将这个操作提到循环外面，缓存在全局缓冲区里，消除了冗余的15次重复计算。

八、哪里调不动：AccelOpt的局限性

公平的评估不仅要看成绩，也要诚实面对局限。研究团队通过细致的观察，总结出AccelOpt遭遇"停滞"的两类典型情况。

第一类是"已经够好了"的停滞。当一个内核在经过若干轮优化后，已经达到了硬件峰值的80%以上，剩下的提升空间自然越来越小。在Matmul(K=5120)这个测试案例中，系统在第7次迭代后达到约82%峰值吞吐量后便不再有显著进展——但这并非因为系统停止了探索，从后续迭代的数据变化可以看出，系统仍在尝试各种不同策略，只是已经没有足够的提升余地了。这类停滞是"胜利后的停滞"，本质上是好事。

第二类是"无从下手"的停滞。BatchMatmul(B=16, K=64)这个案例是最典型的例子。这个算子的问题规模很小，所有数据都能放进芯片的片上缓存中，因此内存传输效率天然就接近100%，没什么可优化的余地。更棘手的是，它的缩减维度K=64只有芯片张量引擎原生支持的维度128的一半，导致硬件利用率很难提上去，当前的NKI接口也缺乏应对这种情况的有效工具。在这个案例的整个优化过程中，几乎所有的性能指标都纹丝不动，系统生成的大多数改动要么正确率极低，要么性能不见起色。这类问题提示了NKI接口本身的设计局限，也指出了未来可以改进的方向。

九、与专家对决：AI能超过人工调优吗

研究团队还做了一个很有说服力的比较实验，把AccelOpt的成果与人类专家手工优化的版本放在一起评判。

对于Mamba序列模型的内核，NKI官方教程提供了三个由人工逐步优化的版本，分别达到28.4%、30.1%和52.7%的峰值吞吐量占比。AccelOpt从同一个起点（28.4%的版本）出发，自主优化到了54.6%，略高于人类专家的最佳版本（52.7%）。更有趣的是，AccelOpt生成的最优内核采用了与人类专家完全不同的循环顺序，走的是另一条路，却到达了更高的终点。

对于RoPE位置编码的内核，NKI官方示例提供了一个达到21.1%峰值的参考版本。AccelOpt从这个版本出发，将其提升到29.6%，实现了1.4倍的加速。

研究团队指出，AccelOpt超越人类专家的根本原因在于**规模效应**：人类工程师通常只能串行地逐个优化内核，而AccelOpt可以同时并行探索大量不同的优化方向。当人类工程师还在仔细思考某个方向是否值得深入时，AccelOpt已经同时跑完了几十个并行实验。

此外，研究团队还把AccelOpt应用到了斯坦福大学CS149并行计算课程的教学中。在2025年秋季学期，他们用AccelOpt优化了一个课程外的二维卷积（Conv2D）内核，从去年的参考实现（9.54%峰值吞吐量）提升到了48.8%。基于AccelOpt提出的优化思路，课程组设计了一道附加题，131支学生队伍中有33.6%成功完成了挑战。学生们在这个过程中具体学到了两个原则：把时序上顺序执行的计算转化为空间上并行执行，以及针对特定工作负载在硬件约束下进行专门化设计。

十、成本控制：聪明地花每一分钱

AccelOpt是一个需要持续调用AI模型接口的系统，因此成本控制是不得不考虑的现实问题。研究团队对各种参数组合和模型选择做了细致的成本收益分析。

在模型选择方面，一个出人意料的发现是：换用不同的规划师模型（无论是小巧的gpt-oss-20b还是强大的Qwen3-235B-Thinking），对最终优化效果几乎没有显著影响；但换用不同的执行师模型，效果差距则非常明显。这说明，在AccelOpt的工作流程中，正确理解并实现代码改动的能力（执行师的任务）远比制定优化方向的能力（规划师的任务）更关键。因此，提升系统性能应该优先考虑用更强的执行师，而不是花钱升级规划师。

在内存配置方面，研究团队发现增大ExpN（扩大记忆库容量、保留更多历史经验）比增大TopK（每轮迭代更激进地更新记忆）更具性价比。以Qwen3-Coder-30B为执行师为例，将ExpN从8增加到16，额外花费12.33美元换来了4.6%的加速提升；而gpt-oss-120b在同样条件下，额外花费13.81美元只带来了0.6%的提升，显示不同模型对扩充历史经验的敏感度差异很大。

在整体对比中，使用Claude Sonnet 4做重复采样的总花费约为5806美元，最终平均加速比约1.222倍；而AccelOpt使用gpt-oss-120b，花费约139美元，达到1.235倍加速。成本差距高达四十倍，效果还略有反超。

十一、不只是Trainium：通用性展望

AccelOpt的设计从一开始就考虑了通用性。束搜索机制和优化记忆库的核心设计与具体硬件无关，要把系统迁移到新的平台，只需要两样东西：一是针对新平台的性能剖析服务，二是包含新平台基础知识的提示词模板。

研究团队已经在GPU平台上做了初步验证。他们将AccelOpt应用于FlashInfer-Bench中的24个Triton内核（运行在英伟达H100 GPU上），使用gpt-oss-120b实现了平均1.27倍的加速，其中对GQA解码内核的峰值加速达到3.19倍。研究团队认为，由于大型语言模型在训练数据中见过大量GPU编程的示例，AccelOpt在GPU上的效果应该比在相对陌生的Trainium上更好，这一初步结果也印证了这个判断。

说到底，AccelOpt做的事情是把"让AI芯片发挥最大效能"这件原本需要顶尖专家耗费数年时光的工作，部分自动化了。它不是万能的——面对某些特殊情况它也会束手无策，也还无法跨越硬件接口本身的设计限制——但它确实证明了一件事：AI可以通过自我探索和经验积累，在一个全新的、陌生的技术领域里，逐步逼近甚至偶尔超越人类专家的水平，而且代价可以比人工低得多。

对于普通用户来说，这意味着未来大模型的推理速度可能更快、运行成本更低，因为驱动这些模型的底层代码会越来越高效。对于AI工程师来说，这意味着他们不再需要在每款新芯片上从头开始漫长的调优之旅，而是可以把这项繁重的初步工作交给系统自动完成，然后在系统探索的基础上做更有创意的工作。

AccelOpt和NKIBench的代码与数据已经开源，这14个测试任务也作为了一个起点，供未来的研究者继续深入探索。对这个方向感兴趣的读者，可以通过arXiv编号2511.15915v2查阅完整论文，以及在GitHub上访问zhang677/AccelOpt项目获取代码。

Q&A

Q1：AccelOpt的优化记忆库和普通AI的上下文窗口有什么区别？

A：普通AI的上下文窗口是单次对话的临时记忆，对话结束后清空。AccelOpt的优化记忆库则是跨轮次持续积累的经验档案，专门存储那些经过实际硬件验证、带来真实速度提升的代码改动规律。它会被整理成通用的优化原则和代码片段对，并在之后的每一轮优化中作为参考提供给规划师智能体，相当于一本通过实战不断扩充的经验手册，而不是每次从零开始。

Q2：AccelOpt能用来优化普通Python或C++代码吗？

A：目前AccelOpt的设计专门针对AI加速器的内核程序（如Trainium的NKI程序和GPU的Triton程序），它的优化逻辑依赖硬件性能剖析数据（如内存带宽占用、引擎利用率等），以及对特定硬件架构的专业知识。不过系统的核心框架——束搜索加上自我更新的经验记忆——理论上可以迁移到其他代码优化场景，只需要替换对应的性能评估工具和领域知识提示词。

Q3：NKIBench中哪类内核最难优化？

A：从测试结果来看，规模较小、数据能完全装入片上缓存的内核最难通过代码层面的优化来提升性能，比如BatchMatmul（B=16, K=64）这个案例。这类问题的瓶颈不在于数据搬运或计算冗余，而在于硬件规格本身的限制——例如Trainium的张量引擎原生支持128维度的缩减，而K=64只有这个维度的一半，导致硬件利用率天花板极低。AccelOpt在这类情况下几乎无能为力，这也提示了NKI编程接口本身还有改进空间。