2025-12-10 17:06:05 中华网
据悉，昇思MindSpore开源社区将于 2025 年 12 月 25日在杭州举办昇思人工智能框架峰会。本次大会，昇思MindSpore AKG团队会展示 AI驱动的算子自动生成相关的技术内容。面向本次大会的昇思人工智能框架技术发展与行业实践论坛，本文基于Agent算子生成技术发展实践就相关领域技术进行深入探讨。
导言
随着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算子的生成、搜索调优等。在大幅减少人参与算子开发的同时，保证了算子生成的正确率与性能。为昇腾平台提供更优的开发体验。
算子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 框架图示
关键设计
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 图示
当前效果展示
4.1 AIKG支持DSL情况
当前AIKG通过文档驱动式接入，已快速打通部分主流 DSL 的流程，部分结果如下:
实验环境：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深度优化，普遍可达成更优的算子性能。
结语
本项目是昇思MindSporeAKG团队在Agent时代的一次尝试，如何将AI编译、自动生成、算子优化等技术结合Agent产生更大的价值，是一件令人激动的事情;长久以来，非CUDA系的软硬件一直受困于编译、算子等底层能力生态支持度缺乏，易用性一直是各大厂商的老大难问题;而Agent+的路线目前看来有望缓解这一困境。
AIKG会持续在这个方向上快速迭代演进，希望在AI民主化、算子民主化的过程中做积极的贡献。
当前 AIKG 已作为 AKG 项目子仓于2025年6月开源(br_aikg分支)，欢迎大家使用、评论及参与共建。代码地址：https://gitee.com/mindspore/akg/tree/br_aikg/
本次在杭州举办的昇思人工智能框架峰会。将会邀请思想领袖、专家学者、企业领军人物及明星开发者等产学研用代表，共探技术发展趋势、分享创新成果与实践经验。欢迎各界精英共赴前沿之约，携手打造开放、协同、可持续的人工智能框架新生态!
