01 导言

随着AI模型快速迭代与AI软硬件不断演进，AI产业对高质量算子的需求愈发强烈：

• 模型侧：大模型（LLM）领域，稠密模型、MOE、MLA、多模态场景众多；稀疏、量化、KVCache压缩等技术又为算子生成带来更多复杂性与多样性。推荐类、 CV类模型同样在持续演进，与LLM截然不同的算子场景同样存在强烈的优化落地需求；
• 硬件侧：随着国内外AI芯片的不断升级和演进，硬件架构、特性、参数各有不同；为提高AI芯片可用性，构建性能优势，定制化算子优化方案需求强烈；

与此同时，基于LLM的算子生成技术逐渐成为业界研究热点。自2025年始，随着大模型的代码生成能力日益完善，用LLM来编写代码已在各类CodeAgent、AI IDE中落地部署；在算子生成领域，各大知名高校、厂商纷纷开始投身到LLM生成算子的探索工作当中。

KernelBench：提供了一套易用的LLM生成算子验证流程基准

昇思MindSpore AKG 团队于2025年初开始对相关问题进行探索，尝试构建一套更为通用的 基于 Agent 的算子生成解决方案，全自动实现多DSL算子的生成、搜索调优等。在大幅减少人参与算子开发的同时，保证了算子生成的正确率与性能。为昇腾平台提供更优的开发体验。

02 算子Agent 框架简介

AKG提供了一款 AI 驱动的算子代码生成工具，利用大语言模型(LLM)的代码生成能力，通过LLM规划和控制多Agent 协同，完成多后端、多类型的AI算子生成和优化。

算子 Agent 框架的主要构成如下：

1、Agents

框架运行围绕着多个Agents展开，包括但不限于：

• Designer：算子设计Agent，生成 Unified Sketch，为Coder提供优化指导；
• Coder：编码Agent，支持多前/后端的编码任务；
• Conductor：负责调度指挥职能，分析Task运行状态，决定执行流程，提出建议；
• Verifier：提供集成工具集，用于编译、运行、精度比对、性能分析等。

2、自定义文档系统

采取文档驱动式接入方案， 通过统一的文档规范与配置声明，使新的 DSL、前端框架与后端硬件在不修改 框架 本体的前提下完成接入，降低耦合与维护成本。

3、搜索框架

在持续性能优化场景中，可通过调用搜索框架，通过多轮迭代与 LLM 自主分析优化方向，持续优化算子代码性能表现。

4、Database with RAG

Database提供高性能算子数据存储、特征提取的功能；配合多层级混合算子检索方案，为各类场景提供更准确的算子示例，提升整体生成的正确性与性能。

AIKG 框架图示

03 关键设计

3.1 基于 Unified Sketch 的代码设计与生成

区别于传统的“文字 --> 算子代码”，算子Agent采用先设计Sketch，再生成代码的思路。让各Agent聚焦关注自身任务问题，进而提升各Agent执行效果。

• Sketch设计: 整合算子特点信息、硬件架构信息、不同算子启发式设计规则等，融合更多优化方案设计，用于提升生成算子效率和性能。
• 代码生成：严格按照API文档、示例代码，确保生成的代码符合规范要求。

Sketch图示

3.2 基于Conductor的工作流框架
预设的固定流程缺乏对决策的把握，尤其是涉及不同难度的算子生成任务、不同的生成阶段，使用同一套流程会极大干扰生成效果。

两种流程设计对比

为解决这类问题，工作流中添加Conductor Agent作为中间调度者，存储完整的生成历史记录。通过Conductor智能调控任务流向，有针对性地分析生成过程并给出下一步建议，结合人工经验和特定的历史记录，引导后续 Agent 正确生成，推动整个生成过程高效运转。

与通用CodeAgent对比，在AIKG框架下适配算子生成场景，针对性地设置任务调度和提示词生成，为算子生成提供规范的分析和进化方案，是正确生成高效率算子的最简洁的框架流程，在繁复的生成任务中指引高效的生成路线。
3.3自定义文档系统

大部分LLM算子生成方案中，仅针对单一前端/后端组合（Torch + Cudac、Torch + Triton-cuda）的算子优化；优化后的成功示例很难无缝迁移其他前后端中去。
而对 AI 算子的泛化支持又会导致以下问题：

• 多样化组合（前端框架/DSL/硬件）的定制化接入成本过高，常需在框架内部进行定制化改造，并存在组合规模爆炸的风险；
• 规范与示例分散，缺少统一约定，难以沉淀和复用知识。

解决方案：文档驱动式接入（ Doc-Driven Integration）
本框架 将前后端的定义从框架中剥离，通过统一的文档规范与配置声明，让前端框架、DSL、AI硬件在不修改框架的前提下完成接入，降低耦合与维护成本。

以统一文档规范方式接入框架

文档规范被划分为：（1）基础说明文档；（2）API文档；（3）优化建议文档；（4）示例参考等；AIKG 会针对特定文档做关键信息检索、长文本压缩等优化以合理利用各类文档。
3.4精度对比与性能采集
本框架 提供 Torch/MindSpore/Numpy 等多个前端的精度、端到端性能验证、报错日志收集功能，新的 DSL 可快速接入 AIKG 框架以完成验证与反馈。在特殊的性能调优功能例如 Triton 的autotune 等，AIKG 也会提供相应的能力支持。
3.5搜索框架
在基础算子生成流程之上，AIKG搜索模块支持在多任务场景下持续迭代优化算子性能。通过AIKG搜索框架方案，结合生成流程、数据库数据、优化要点，迭代生成高性能算子实现。搜索框架基础设计如下：

• 通过Sketch抽取算法框架和优化重点，兼容不同DSL；
• 检查待优化列表，让大模型重点关注优化侧重点；
• 通过重新抽取Sketch的方式，将算法逻辑和性能数据存储至 Database，后续生成时依据数据库的历史记录提供更合理的优化方案；
• 从本地 Database 中抽取生成历史，作为下一轮生成的启发信息，检查 checklist 中未实现的部分，建议下一轮着重实现；
• 由基础搜索框架扩展，AIKG进一步提供了Evolve类并行多轮搜索功能。

多轮搜索流程图

3.6多层级混合算子检索方案
随着算子数据的不断增加，在算子数据库中的数据条会越来越多，如何高效高质量的利用这些已有信息是进一步提升AIKG效率、性能表现的关键。基于这个认识，AIKG集成了相对应的RAG能力以及解决方案。
而在Kernel代码相似性检索这一任务中，我们识别到Kernel代码由于其特殊性：不同Kernel存在大量的相似代码行，简单的Shape、Tiling变化会导致实际含义差别巨大，而这些信息很难被传统的相似性检索方案识别出来。
为此我们提出一种多层级混合算子检索方案，通过这种更贴近算子逻辑的分析+检索，保证检索结果的大部分内容与当前任务相关。基本流程如下：

• 特征提取
  1）初始任务特征：算子的类型，计算逻辑，shape大小等；
  2）算子代码特征：tiling, reordering, vectorization, layout, pipeline等；
• 分层混合检索
  1）基于后端架构、DSL、算子大类等基础要求在Database内进行粗筛；
  2）将LLM生成的特征转换成Embedding，建立向量索引并逐层筛选。

AIKG RAG 图示

04 当前效果展示

4.1 AIKG支持DSL情况
当前AIKG通过文档驱动式接入，已快速打通部分主流 DSL 的流程，部分结果如下:

DSL	KernelBench Level 1 Pass@4
Triton-ascend	77/100
Cpp	63/100
Tilelang	44/100

实验环境：ASCEND：Ascend Atlas 800 A2；；CPP：ARM64。

注：Triton-ascend的当前对 img2col API缺少支持，30+个卷积算子生成存在难度。除此之外大部分场景都能支持。

4.2 Triton-ascend 性能验证

Agent 在各类型算子中完成基于多轮搜索的Triton KernelBench Level 1性能优化，验证了AIKG 在性能调优方面的能力。同时在Inductor-Triton场景中，通过AIKG深度优化，普遍可达成更优的算子性能。

05 结语

本项目是昇思MindSporeAKG团队在Agent时代的一次尝试，如何将AI编译、自动生成、算子优化等技术结合Agent产生更大的价值，是一件令人激动的事情；长久以来，非CUDA系的软硬件一直受困于编译、算子等底层能力生态支持度缺乏，易用性一直是各大厂商的老大难问题；而Agent+的路线目前看来有望缓解这一困境。

AIKG会持续在这个方向上快速迭代演进，希望在AI民主化、算子民主化的过程中做积极的贡献。
当前 AIKG 已作为 AKG 项目子仓于2025年6月开源（br_aikg分支），欢迎大家使用、评论及参与共建。代码地址：https://gitee.com/mindspore/akg