AKG Agent：大模型驱动的算子自动生成技术探索

在AI模型与软硬件快速演进的今天，高质量算子已成为制约AI性能提升的关键瓶颈。我们看到一个有趣的趋势：算子自动生成技术正从"学术研究"走向"工程实用"。特别是2025年以来，随着大模型代码生成能力的显著提升，业界开始系统性地探索如何将 LLM 应用于高性能算子开发这一垂直领域。各大高校、AI厂商纷纷投入研发，期望通过智能化手段解决算子开发的效率和可移植性问题。

AKG 团队自2025年初起持续投入这一方向，并取得了一系列成果（昇思MindSpore AKG大模型驱动算子生成）。与此同时，我们的思路也在不断演进——与其为单一具体问题打造一个专用工具，不如构建一个通用的多Agent协作框架？从初期的只聚焦算子生成，到如今的 AKG Agent 定位为通用AI Infra多Agent协作平台，我们希望沉淀下来的经验能够服务更广泛的场景：算子迁移、性能调优、子图分析、算子拆分与融合……

01AKG Agent 框架简介

AKG Agent 是一个面向 AI Infra 与高性能计算场景的 LLM 多 Agent 协作框架，通过智能 Agent 协同，将高性能代码的开发与优化效率提升到新的高度。

1.1设计理念

AKG Agent 采用了通用的架构设计：构建一个可复用的多 Agent 协作基础框架，让不同领域的 AI 工作都能受益于智能化的代码生成与优化。

在这个设计理念下，AKG Agent 具备以下特点：

通用性：核心框架不绑定具体场景，可扩展至算子迁移、性能调优、子图分析与算子融合等更多领域

可组合性：Agent、Skill、Tools 可灵活组合，适应不同任务需求

可观测性：完整的Trace追踪系统，让AI决策过程透明可控

可扩展性：通过Skill系统和Agent注册机制，社区可轻松扩展能力

1.2核心架构

AKG Agent 提供了完整的 Agent 基础设施，整体架构如下：

02关键设计

2.1Skill 系统：让 Agent 学会"按需取用"

早期方案把所有知识文档一股脑塞给 LLM，生成简单算子和复杂算子喂的是同一份“大部头”——就像考每一科都带着全部课本，信息过载反而找不到重点。

AKG Agent 的 Skill 系统改变了这一做法。我们将领域知识拆解为一个个独立的 Skill 单元（以 SKILL.md 格式存储），每个 Skill 聚焦于一个特定主题——比如"triton-ascend-basics"。当任务到来时，系统先根据元数据（目标硬件、DSL 类型等）快速缩小候选范围，再由 LLM 根据当前任务的具体特点，智能挑选最相关的少数几个 Skill 注入到上下文中。

这带来了显著的好处：

Prompt 长度大幅缩短，Agent 能够更专注于当前任务真正需要的知识；

不同场景、不同硬件自动加载不同的知识组合；

成功经验可以沉淀为新 Skill ，社区开发者也能通过提交 SKILL.md 轻松贡献知识。

2.2Trace 系统：给 AI 决策装上"行车记录仪"

AKG Agent 的 Trace 系统用一棵树状结构记录了 AI 的完整决策历程。每当 React Agent 做出一个决策——选择了哪个工具、收到了什么反馈——都会作为一个节点挂在这棵树上。与线性日志不同，树状结构天然支持"分叉"：当某一步有多个可能的方向时，你可以在同一个节点上探索不同的分支，就像游戏里的存档读档。

这就意味着：

开发者可以随时回溯到任意历史节点，查看当时的完整状态（包括代码快照、Agent 的输入输出等）

从某个节点重新出发，尝试不同的策略。

程序中途崩溃也不要紧，所有状态都持久化在文件系统中，恢复后可以从上次中断的地方继续。

对于复杂算子的生成与优化，这种"时光机"般的能力让调试效率有了质的提升。

2.3Agent 与工作流：灵活性与稳定性的平衡

用 LLM 来驱动复杂任务，有一个永恒的矛盾：你既希望 AI 足够灵活，能根据情况自主判断下一步做什么；又希望关键流程足够稳定，不会因为模型的一次"脑洞大开"而跑偏。AKG Agent 通过两套互补的执行机制来解决这个问题。

ReAct Agent 代表的是灵活的一面。

实现了经典的 “推理-行动”循环：观察状态 → 推理下一步 → 选择并执行工具 → 根据结果决定后续动作。整个过程由 LLM 自主驱动，没有预设路径。

适用场景：探索性、开放性的任务。比如在 OpenCode 工作台中，用户说“帮我优化这个模型的性能”，ReAct Agent 可以自主决定：先做环境检测 → 再分析瓶颈 → 最后选择优化策略——每一步都灵活应变。

基于 LangGraph 的工作流 则代表稳定的一面。

将多 Agent 协作的流程预先编排为一张有向图：哪个 Agent 先执行、结果传给谁、失败了怎么重试、最多重试几次——这些都是确定性的。比如算子生成的 Default 工作流：Designer → Coder → Verifier 的流转路径是固定的，不会因为 LLM 的随机性而跳过验证环节。路由函数还内置了步数限制和防循环机制，确保流程不会失控。

💡 一句话总结 工作流：适合有明确范式的任务（如算子生成、代码迁移），保证流程的可靠性和可复现性；

ReAct Agent 适合需要自主决策的场景（如交互式问题排查、开放式优化探索），充分发挥 LLM 的推理能力。

两者在 AKG Agent 中并不互斥——实际上，工作流中的每个节点本身就可以是一个 ReAct Agent，在预定义的流程框架内保留局部的灵活性。

03落地场景：AI 算子自动生成

AKG Agent 的首个落地场景是 AI 算子代码生成——这正是原先项目的核心能力。在AKG Agent 框架下，我们重构了算子生成流程，使其更加清晰、可维护、可扩展。

3.1支持的后端与 DSL

3.2工作流

算子生产Agent 围绕算子生成的不同阶段需求，提供了三种预定义工作流。

Default 工作流是最完整的流程，走"设计→生成→验证"的全链路——先由 Designer 分析需求并生成高层次的算法草图，再由 Coder 将草图转化为具体 DSL 代码，最后由 Verifier 编译运行并与参考实现对比验证，验证失败时 Conductor 会分析错误并引导 Coder 自动修正，如此迭代直到通过。

KernelGenOnly 工作流则更加轻量，跳过独立的设计阶段，由 KernelGen 直接基于 Skill 知识生成代码后进入验证-修复循环，适合模式较为明确、不需要复杂设计推理的算子。

VerifierOnly 工作流不涉及任何代码生成，仅对用户传入的已有代码执行正确性验证与性能测量，适合手写算子的快速检验或性能基准测试。

3.3算子优化策略

3.3.1Evolve 算法详解

Evolve 采用进化算法思想，通过多轮迭代和变异操作，持续探索更优的算子实现：

核心特点：

同步迭代：每一轮需要等待所有并行任务全部完成后才进入下一轮选择-变异，轮间存在等待开销

岛屿隔离模型：将种群划分为多个独立岛屿，各岛屿独立进化并定期迁移精英个体，保持种群多样性，避免过早收敛

精英池机制：维护全局精英池，从多个岛屿汇聚最优个体，变异来源更丰富

3.3.2Adaptive Search 算法详解

Adaptive Search 采用UCB（Upper Confidence Bound）策略，实现自适应的异步流水线搜索：

核心优势：

即时递补：任意任务一完成，立即基于当前最优父代启动新任务，无需等待同批次其他任务结束，最大化资源利用率

UCB 智能选择：每次扩展时基于 UCB（Upper Confidence Bound）策略选择最有潜力的父代进化，自动平衡探索与利用，趋向极致性能

收敛更快：持续滚动式搜索，相同时间内可完成更多轮迭代，更快找到高性能实现

04OpenCode 集成：从工具到工作台

AKG Agent 与 OpenCode 深度集成，构建了一个完整的 AI Infra 智能工作台。OpenCode 通过内置的 Agent 和 Skills 机制，实现了对 AKG Agent 的无缝调用。用户在 OpenCode 工作台中，只需用自然语言描述需求，OpenCode 会自动识别任务类型，并通过 akg_cli 将算子相关的专有任务委托给 AKG Agent 框架处理。整个过程对用户透明——环境检测、任务构建、工作流调度、结果验证都在后台自动完成。

这意味着开发者有了两种互补的使用方式：akg_cli 命令行直接调用AKG Agent 执行算子相关任务；OpenCode 工作台则提供交互式的对话体验，间歇性调用AKG能力，适合探索性的优化任务和快速原型验证。两种入口调用的是同一套 AKG Agent 框架，能力完全一致，开发者可以根据场景自由选择。

05AKG CLI：开箱即用

AKG Agent 提供了 akg_cli 命令行工具，让开发者可以快速上手。

具体安装配置：

# 1. 创建环境
conda create -n akg_agents python=3.11
conda activate akg_agents
# 2. 克隆仓库
git clone https://gitcode.com/mindspore/akg.git -b br_agents
cd akg
# 3. 安装依赖
pip install -r akg_agents/requirements.txt
pip install -e ./akg_agents --no-build-isolation
# 4. 快速配置llm
export AKG_AGENTS_BASE_URL="https://api.deepseek.com/beta/"
export AKG_AGENTS_API_KEY="your key"
export AKG_AGENTS_MODEL_NAME="deepseek-chat"
export AKG_AGENTS_MODEL_ENABLE_THINK="disabled"

实机演示：

06结语

AKG Agent 项目是 AKG 团队在 Agent 时代的一次重要尝试。从 2025 到 2026，我们的思路也在不断演进：

我们之前专注于解决算子生成这一具体问题，探索了 Unified Sketch、Conductor、文档驱动式接入等关键技术方案，取得了不错的成果。而现在，我们将这些经验沉淀为通用的多 Agent 协作框架，希望让更多 AI Infra 相关场景都能受益于智能化的代码生成与优化。

无论是哪个硬件平台，快速迭代、高效开发与持续优化算子始终是刚需。随着模型架构日新月异、硬件能力不断演进，算子开发的效率和质量直接决定了AI系统的落地速度。而 Agent + 通用框架的路线，目前看来有望从根本上提升这一环节的生产力。

AKG Agent 会持续在这个方向上快速迭代演进，希望在 AI 民主化、算子民主化的过程中做积极的贡献。

07实战赛事邀请

当前 AKG Agent 已作为 AKG 项目子仓开源（br_agents分支），相关最新进展也将会在 AKG SIG 例会、知乎专栏上进行分享，欢迎大家使用、评论及参与共建。

项目地址：

https://gitcode.com/mindspore/akg/tree/br_agents

相关文档：项目提供了详细的架构、Agent 体系、Skill 系统、工作流等文档。

纸上得来终觉浅，想要亲手验证 AKG Agent 的硬核实力？不妨报名参加第八届 CCF 开源创新大赛，用这套技术栈冲击 10万奖金池！