本文作者Jaber Jaber和Osama Jaber来自Rightnow AI，该研究成果（预印本）以《AutoKernel: Autonomous GPU Kernel Optimization via Iterative Agent-Driven Search》为题发布于arXiv平台，发布日期为2026年3月22日。这项研究属于计算机科学中的高性能计算与机器学习系统领域，具体聚焦于利用大型语言模型驱动的自动化方法来优化GPU核心代码。
本研究的学术背景源于现代大规模Transformer模型运行时性能高度依赖于少数几个关键计算核心。虽然NVIDIA cuBLAS等厂商库为常见操作提供了高度优化的实现，但深度学习架构的创新速度往往快于官方库的更新支持，导致新型算子（如分组查询注意力、RMSNorm等）在发布初期缺乏高性能实现。手工优化GPU核心是一项极度耗时且需要深厚专业知识的工作，涉及内存布局、平铺策略、指令集选择等诸多复杂微架构层面的权衡。尽管已有工作探索使用LLM生成GPU核心，但一次性生成的成功率较低。因此，本研究旨在开发一个能够自主、迭代地优化GPU核心的端到端系统，将原本需要数周专家工作的流程转化为一个无人值守的自动化过程。
研究的具体目标是开发一个名为AutoKernel的自动化框架。该框架的核心思想是将专家的工作流程（编写候选核心、测试性能、保留改进、丢弃性能回退、重复）进行机械化。给定一个PyTorch模型，AutoKernel能够通过性能剖析识别计算瓶颈，根据阿姆达尔定律确定优化优先级，并通过一个自主运行的智能体在数百次实验中迭代改进Triton或CUDA C++核心实现，全程无需人工干预。
该研究的工作流程十分详细，主要分为三个核心阶段，构成一个完整的自动化闭环。整个系统由超过9200行Python代码、一个909行的智能体指令文档、涵盖两个后端（Triton和CUDA C++）的18个起始核心实现、一个六层优化手册以及一个集成验证套件组成。
第一阶段是模型剖析与核心提取（阶段A）。系统首先接受一个PyTorch模型，使用Torch Profiler进行性能剖析，记录每个CUDA核心在可配置迭代次数下的GPU时间消耗。随后，系统通过核心名称的模式匹配将内核分类到九种支持的操作类型，例如矩阵乘法、FlashAttention、RMSNorm等。剖析器能够检测GPU硬件类型，并针对未知GPU估算其理论峰值性能。接着，提取器会根据剖析结果，筛选出支持的操作类型，并为每个目标核心生成独立的、包含模型特定形状信息、计算量公式和数据类型容差设置的起始实现文件。最关键的一步是，系统会应用阿姆达尔定律来计算每个瓶颈核心的优化潜力，并据此制定优化计划，将努力优先分配到对整个模型端到端性能影响最大的核心上。
第二阶段是智能体驱动的优化循环（阶段B）。这是AutoKernel的核心创新。整个循环遵循算法1，可概括为“编辑-测试-保留/回退”。在每次迭代中，智能体根据一份详细的、长达909行的指令文档（program.md）来修改单个核心文件。这份指令文档实际上是一个“六层优化手册”，编码了专家知识：第1层是块大小调优，第2层关注内存访问优化（如合并加载、软件预取），第3层是计算优化（如TF32累加、循环不变外提），第4层是高级技术（如Split-K、持久化核心），第5层是架构特定优化（如Hopper上的TMA），第6层则是特定核心的专有优化策略。智能体修改后，系统会执行一个严格的五阶段正确性验证（图2）。这个验证流水线包括：1）冒烟测试（快速检查编译和基本数值）；2）形状扫描（在10多种配置和三种数据类型上运行）；3）数值稳定性测试（使用对抗性输入，如极大值或极小值）；4）确定性验证（三次运行，要求逐位相同）；5）边缘情况测试（非2的幂次维度）。只有通过所有五阶段验证的候选核心才会进入性能基准测试。如果候选核心性能相比最佳版本提升超过1%，则被保留，否则通过Git回退到上一个版本。每次迭代大约耗时90秒，每小时可进行约40次实验。一个优化会话（例如通宵运行）可以产生300至400次实验。整个循环由一个编排器管理，它会根据“继续优化”判据（如连续5次回退、达到GPU峰值性能的90%、超时2小时或达到2倍加速）来决定何时停止对当前核心的优化，并转向下一个高优先级核心。
第三阶段是端到端验证与社区集成（阶段C）。在完成优化循环后，系统会进行端到端模型验证，确保优化后的核心不仅自身性能提升，并且集成回模型后能正确运行并带来预期加速。此外，研究还展示了AutoKernel与现有基准测试套件（如KernelBench）的集成，以及将优化后的核心导出至HuggingFace平台分发的工具链。
该研究的主要结果丰富且具有说服力。在配备NVIDIA H100 GPU的实验环境中，研究人员评估了AutoKernel生成的Triton核心的性能，并与两个基线进行了对比：PyTorch Eager执行（使用cuBLAS和Aten原语）以及使用torch.compile（max-autotune模式）编译的模型。结果显示，在大多数测试配置中，AutoKernel的Triton核心性能显著优于基线。尤其对于内存密集型核心，性能提升最为突出：在8192x8192的尺寸下，RMSNorm核心相较于Eager和torch.compile分别实现了5.29倍和2.83倍的加速，吞吐量达到2788 GB/s，约为H100峰值带宽的83%。Softmax和Cross-Entropy也分别实现了2.82倍/3.44倍和2.21倍/2.94倍的加速。对于计算密集型的矩阵乘法，虽然起始实现（278 TFLOPS）距离高度优化的cuBLAS（800+ TFLOPS）仍有差距，但在2048^3尺寸上仍能击败torch.compile 1.55倍。所有34个测试配置都成功通过了五阶段正确性验证，无一失败。
研究还提供了社区部署的实际案例，进一步证明了系统的实用价值。在一个名为vectorsum_v2的B200 GPU基准测试中，经过AutoKernel优化迭代的Triton归约核心获得了排行榜第一名的成绩。另一个案例显示，社区用户通过一次约3分钟的智能体交互提示，生成了一个FP4精度的矩阵乘法Triton核心，在多个形状上其性能超越了高度手工优化的CUDA模板库CUTLASS，速度达到后者的1.63到2.15倍。这些结果有力地说明了AutoKernel不仅是一个研究原型，更是一个能产生实际高性能代码的工具。
本研究的结论是，AutoKernel成功地通过一个迭代的、智能体驱动的闭环，实现了GPU核心优化的自动化。它将一个需要数周专家时间的活动转变为一个可以通宵自主运行的流程。系统的核心贡献在于将模型级剖析与阿姆达尔定律调度相结合，确保了优化工作在最重要的地方展开；其严格的五阶段正确性验证机制保障了优化的可靠性；对Triton和CUDA C++双后端的支持兼顾了迭代速度与控制深度；而编码了专家知识的六层优化手册则为智能体提供了有效的指导。
该研究的亮点和创新之处主要体现在以下几个方面：首先，方法论的简洁与有效性：它没有采用复杂的多智能体架构，而是将专家工作流提炼为一个简单的“编辑-测试-保留/回退”循环，并辅以固定且严格正确的评估体系，避免了智能体“博弈”评估标准的风险。其次，端到端与优先级驱动：不同于以往孤立优化单个核心的工作，AutoKernel从完整的PyTorch模型出发，通过剖析和阿姆达尔定律来优先优化对整体性能影响最大的瓶颈，使得优化成果能直接转化为端到端的应用加速。第三，严谨的工程保障：五阶段正确性验证流程是系统可靠性的基石，它捕捉了从编译错误到数值稳定性、非确定性等各类潜在问题，确保了所有性能提升都是在保证绝对正确性的前提下取得的。第四，双后端设计与社区集成：同时支持Triton（快速原型）和CUDA C++（深度优化），为不同场景和优化目标提供了灵活性。与KernelBench的集成和向HuggingFace平台的导出能力，则提升了其学术评估和实际部署的便利性。
最后，研究也坦诚了其局限性，例如其能力受限于底层LLM的代码生成水平，过于复杂的优化策略（如软件流水线、自定义PTX发射）可能超出当前智能体的能力范围。系统目前也主要针对单GPU的独立核心优化，分布式核心和多设备内存管理尚在范围之外。未来可探索的方向包括基于种群的跨多GPU实例搜索、利用历史实验数据训练搜索策略，以及基于性能剖析计数器引导的突变等。
AutoKernel是一个设计精良、工程实现扎实且效果显著的自动化GPU核心优化系统。它不仅为学术界提供了一个可复现、可扩展的研究平台，也为工业界提供了一种能够显著降低高性能计算开发门槛的实用工具，在深度学习系统自动化优化的方向上迈出了坚实的一步。