用AI进行GPU编程，GEAK-Agent让小白也能上手

始智AI wisemodel.cn社区是源自中国的中立开放的AI开源社区。正在招募 | 新一期开源共创志愿者计划，欢迎加入共同成长。wisemodel推出邀请注册奖励活动，最高可得算力券+token包380元奖励，欢迎参与和支持！

AI 任务越来越复杂，硬件也越来越多样，现在大家都急需一种智能系统，不用手动调试，就能直接生成高性能 GPU 内核。不管是运行顶尖大语言模型的超大规模数据中心，还是学术研究用的高性能计算集群，软件和硬件得高效配合才能落地，这种需求就更迫切了。

虽然现在大语言模型在代码生成上挺厉害，但要写出能正确运行、还高效的 GPU 代码，依然是个大难题。而 AI 辅助开发 GPU 内核（尤其是针对 AMD Instinct MI300X 这类新型硬件），能大大提升开发效率，加快部署速度。

AMD 推出了 GEAK（生成高效 AI 中心 GPU 内核），一款专门针对 AMD Instinct GPU 的自动 Triton 内核生成智能体框架。它把最先进的大语言模型和结构化推理、反馈循环结合起来，就是要把 AI 辅助代码生成的能力再提一个档次。

同时还发布了两个评估基准：一个是 TritonBench 修订版（从 TritonBench-G 改编而来，包含 184 个内核，测试工具更严格），另一个是全新的 ROCm Triton 基准（从 AMD 开源 ROCm 仓库挑选 30 个真实场景内核），用来测试 AI 生成内核的正确性和运行性能。

实际测试结果很亮眼：GEAK 比直接用大语言模型生成代码强太多 ——TritonBench 修订版上正确生成率高达 54.89%，ROCm Triton 基准上更是达到 63.33%，而直接提示大语言模型不做智能体反馈的话，正确率还不到 15%。另外，GEAK 生成的内核在 TritonBench 修订版上，比原来的参考版本平均快 2.59 倍。还对 ROCm Triton 基准中的一个高性能内核做了详细研究，并且已经开源了 GEAK 智能体和评估框架。

目前，GEAK-agent已开源上线始智AI-wisemodel开源社区，8月底举办的AMD AI 速训营24小时 GPU黑客松的一个赛题就是基于GEAK-Agent设计的，几十个团队已经在参赛时深度体验和尝试，欢迎更多感兴趣的朋友前往体验和尝试。

模型地址

https://wisemodel.cn/codes/motata/GEAK-agent

01.

挑战与突破

随着程序合成技术的发展，大家对“严格且多维度”的基准测试需求越来越迫切。现在大部分编码基准测试都只盯着“功能对不对”，比如HumanEval是专家写的编程任务加配套测试用例，MBPP是靠众包收集问题，都是用手工设计的测试或独立执行环境来评估。

为了覆盖更多场景、提升扩展性，近几年开始流行“自动生成测试用例”方法，这也让它在软件工程、开发者辅助这些领域用得更广泛了。  但除了“能不能用”，“跑得够不够快”也成了基准测试的关键维度，越来越受重视。可现有框架大多是针对“竞争式任务”或“顺序执行代码”，对并行编程、硬件相关的评估支持特别少。

虽然有少数基准测试针对CPU并行编程，但GPU相关的寥寥无几——要知道现在部署大规模深度学习模型，硬件效率可是越来越重要了。

为了补上这个缺口，KernelBench和TritonBench这类基准测试就应运而生了，专门评估大语言模型生成GPU内核的能力。其中KernelBench重点看不同任务的正确性和运行效率，还搞了“fast_p”这类指标；TritonBench则聚焦大语言模型用Triton领域特定语言（DSL，简单说就是专门针对某类场景设计的语言）生成高性能代码的难点。

这些测试都暴露了一个问题：普通代码生成和“高性能、懂硬件”的代码生成差距太大，必须要有专门的、面向计算的评估框架。

用于代码生成的 LLM

近几年很多研究都在试让大模型写代码，像DeepSeekCoder、Qwen-Coder这些模型在通用编码基准测试里表现不错，但遇到“性能敏感的底层代码”就不行了。

也有专门为“内核生成”训练的模型，比如KernelLLM，但它只能把PyTorch代码转换成Triton内核，通用性太差。而我们的智能体不一样，既能靠自然语言指令生成Triton代码，也能参考已有代码来写，灵活性更高，还能结合上下文。另外，KernelLLM这类专用模型，没有GPT-4.1、Claude 4这些顶尖大模型的强大推理和问题解决能力，而处理复杂优化任务时，往往要权衡各种约束，这些能力恰恰是关键。

推理时计算扩展

“推理时计算扩展”简单说就是大模型在“干活（推理）的时候”，多分配点计算资源（比如给更长的上下文、多跑几次推理、并行生成），不用额外训练就能提升复杂任务的表现。  比如Chain-of-Thought提示法，就是让模型“一步步思考”来增强推理能力；Reflexion框架则是把反馈记下来，指导后续决策，让智能体越用越聪明。

GEAK智能体也用了这种思路，还加了类似Reflexion的调试模块，能根据之前的执行反馈，一次次优化生成的代码。

02.

GEAK技术介绍

如上图所示，该智能AI系统包含生成器、反思器、评估器和优化器四个核心模块：生成器根据用户查询与上下文信息生成代码；评估器采用级联设计，先通过功能测试验证代码正确性，未通过则将错误跟踪反馈给反思器深入分析，通过后继续评估延迟、内存效率等性能指标；反思器会针对执行不正确的代码，结合代码本身与错误跟踪定位潜在问题；优化器则以功能正确的代码为输入，制定策略提升其延迟表现与计算效率，同时系统还集成了多种可组合、可配置的技术以进一步提升智能体性能。

单样本提示（1-shot prompting）

为了实现单样本提示，会从现有数据集中选取最相似的 Triton 代码放入提示中。用于单样本提示的数据集与基准测试集不重叠。我们发现，从数据集中检索单样本时，基于代码相似性比基于指令相似性更有效。

知识注入（Knowledge Injection）

在提示中加入了关于编写高效 Triton 内核的领域特定知识，包括详细的硬件规格。这种底层优化原则的融入显著提升了大型语言模型（LLM）生成代码的准确性和质量。

反思机制（Reflexion）

为了增强系统的自我纠正能力，引入了Reflexion[13]模块，该模块支持内省式调试和迭代改进。当 LLM 生成的内核未通过功能测试时，产生的错误跟踪会作为反馈提供给反思器，以进行进一步分析和纠正。智能体会负责分析失败原因并提出相应的解决方案，通过增加迭代次数来改进内核生成。

内核智能体的 LLM 选择

尝试了多种 LLM，如 GPT-4.1、O1 和 Gemini 2.5 Pro，观察到不同模型产生的输出存在显著差异，这表明底层 LLM 的能力会对结果产生重大影响。

作为优化器的 LLM

优化器 LLM 的任务是根据之前的代码生成结果及其相应的性能指标，确定潜在的优化方向。这些历史记录（包括生成的代码和相关的性能结果）会按性能升序排列。这种结构化的呈现方式有助于引导 LLM 提出更有效的优化策略[14]。

调试陷阱（Debugging trap）

当 LLM 生成的代码存在错误时，错误跟踪会提供给反思器进行纠正。然而，发现有时代码经过多次反思循环后，仍然存在相同的错误，这就是所说的“调试陷阱”。

为了防止智能体陷入调试陷阱，通过参数max_perf_debug_num对每个代码片段的调试尝试次数进行限制。如果代码在达到该阈值后仍然失败，智能体会被要求放弃当前方法，生成新的策略和全新的代码。

并行缩放（Parallel Scaling）

为了确保可靠的内核生成，会并行且独立地运行多个 GEAK 实例多次。为了增加生成代码的多样性，LLM 输出会通过将温度参数设置为 1 来进行采样。

实验表明，这种生成多样性可以产生原本未被充分探索的正确且更快的内核。此外，进一步的实验表明，结合顺序缩放和并行缩放可以带来额外的性能提升。为了获得无偏的准确率估计，使用 pass@k 指标进行评估。

03.

基准测试

团队实验了 Triton 内核评估基准测试，包括从 TritonBench-G 修改的 184 个内核集和从各种开源 ROCm 存储库中选取的 30 个内核集。

TritonBench-revised 基准测试

Triton 内核和相应的评估代码库借鉴并改编自 TritonBench，但做了以下增强：TritonBench-G 的184个内核中有37个在AMD GPU上运行时有错误（如共享内存错误、无效的 HIP 参数（HIP是AMD的异构计算接口，类似于NVIDIA的CUDA）、ModuleNotFound），修复了这些内核以兼容 AMD GPU。

对于一些内核，原始 TritonBench 存储库缺少调用函数，导致执行准确性脚本实际上比较的是空字符串输出。这一问题通过以下方式得到纠正：为缺少调用函数的内核添加测试函数调用。开发基于容差的测试函数输出张量比较，代替 STDOUT 字符串比较。确保单元测试中用于生成输入的随机张量生成具有一致的种子。对于一些内核，单元测试编写不一致，导致输出不符合预期，修复了这些内核，使其具有一致的单元测试。

ROCm Triton 基准测试

以下内核由 Triton 专家工程师编写并公开，以帮助围绕AMD GPU高效运行AI工作负载构建生态系统，借助前沿LLM重构这些内核及其相应单元测试，使其格式与TritonBench-revised基准测试一致。这样做能够以一致的方式比较智能体的准确性和效率。

评估指标

团队采用了与 TritonBench 一致的评估指标：调用准确性指 AI 生成内核可成功编译且无错误运行的比例；执行准确性仅针对编译成功的内核计算，即满足所有单元测试的内核占比；加速比则仅针对通过全部单元测试的内核，定义为所有单元测试中参考基准内核延迟中位数与生成内核延迟中位数之比的期望值，用于衡量生成内核相对参考内核的执行时间改进效果。

04.

实验结果概要

以 GPT-4.1、Gemini 2.5 Pro 等前沿 LLM 直接提示生成内核作为基准，其执行准确率低、性能不佳，非 CUDA 平台（如 ROCm）挑战尤为突出；引入单样本提示后，准确率和部分场景加速比有所提升，但存在 Gemini 2.5 Pro 加速比下降的例外情况。

GEAK 系统表现显著优于直接 LLM 提示：在 MI300 上，TritonBench-revised 基准测试中实现 54.89% 执行准确率和 2.59 倍加速比，ROCm 基准测试中达 63.33% 执行准确率和 0.92 倍加速比。

研究显示，GEAK 生成的triton_test_flip.py内核比专家编写版本快2.26倍，核心优势在于优化内存访问模式（单遍操作、合并访问）、降低寄存器压力、提升缓存效率及显式处理边界条件。

消融实验结果显示：顺序缩放方面，增加迭代次数可使GEAK的准确率和性能稳步提升，19次迭代时调用准确率达63.04%、执行准确率44.02%，加速比长期保持在1.4倍以上；并行缩放方面，固定10次顺序迭代后，并行运行次数的增加使准确率呈近乎对数线性增长，与顺序改进形成互补，且在MI250、MI300上均展现出良好的鲁棒性和可移植性；模块影响方面，知识注入、单样本提示、优化器三大模块均发挥积极作用，其中知识注入显著提升准确率，单样本提示进一步优化正确性，优化器则大幅提升加速比，三者结合时执行准确率达40.76%，加速比达1.45倍。

ROCm 基准测试中的推理时间计算扩展性研究

TritonBench 修改基准测试的推理时间计算扩展性研究