AKG Kernel Agent — LLM 驱动的昇腾 NPU 算子自动生成系统

2026 MindSpore 开源大赛 · 赛道二 · 团队「崇理」

项目概述

基于 LLM Agent 自动将算子自然语言描述编译为可在 Ascend 910B NPU 上运行的 AscendC 代码。系统串联语义解析→知识检索→签名设计→代码生成→编译验证→正确性/性能验证的完整闭环，支持单算子推理和 Benchmark 批量评估。

环境要求

组件	版本	说明
Python	≥ 3.10	Agent 运行时
CANN	8.5.0	昇腾计算平台
NPU	Ascend 910B	不可用 310 开发板
CCE 编译器	bisheng (CANN 内置)	`--npu-soc=Ascend910B3`
CUDA	不需要	仅 NPU 环境

依赖安装

pip install -r requirements.txt

核心依赖：

numpy — 正确性验证 (numpy 参考实现)
pyyaml — 配置文件解析
chromadb — 向量知识库 (RAG)
langgraph — Agent 工作流编排
sentence-transformers — Embedding 模型

可选 (LLM API)：

openai — 调用远程 LLM
transformers — 本地模型推理

环境变量

export CANN_PATH=/home/ma-user/Ascend/cann-8.5.0    # CANN 安装路径
export ASCEND_SOC=Ascend910B3                        # NPU 型号
export LLM_API_KEY=your_api_key                      # LLM API (DeepSeek/OpenAI 等)
export LLM_BASE_URL=https://api.deepseek.com/v1      # LLM API 地址

快速开始

1. 构建知识库

# 将 AKG/CANN 文档放入 rag/knowledge_base/
# 构建向量索引
python main.py --build-index --doc-dir rag/knowledge_base

2. 单算子生成

# 生成 float16 GELU 算子
python main.py --op "实现 float16 类型的 GELU 激活函数，shape [B, N, D]"

# 指定输出目录
python main.py --op "SiLU 激活函数" --output ./output/silu

3. Benchmark 批量评估

# 对整个 Benchmark 运行生成 + 评估
python main.py --benchmark benchmark/benchmark.json --output ./reports

# 输出报告: reports/benchmark_report.json

项目结构

akg-kernel-agent/
├── main.py                    # 主入口 (argparse, 单算子/批量模式)
├── requirements.txt           # Python 依赖
├── config/
│   └── default.yaml           # LLM/NPU/RAG 配置
├── agent/                     # Agent 层
│   ├── workflow.py            # LangGraph 六步工作流编排
│   ├── state.py               # AgentState 共享状态定义
│   └── agents/                # 五个专项 Agent
│       ├── semantic_parser.py   # 自然语言→算子元信息
│       ├── code_generator.py    # CoT + 模板代码生成
│       ├── signature_designer.py# 输入输出签名设计
│       ├── compiler.py          # CCE 编译验证 + 修复
│       └── verifier.py          # numpy 正确性 + 性能预估
├── rag/                       # 检索增强生成层
│   ├── indexer.py             # ChromaDB 索引构建
│   └── retriever.py           # 向量+BM25 混合检索
├── benchmark/                 # Benchmark 框架
│   ├── benchmark.json         # 20 个测试算子定义
│   ├── loader.py              # Benchmark 加载器
│   ├── evaluator.py           # 指标计算 (Pass@1, Pass@3, 性能)
│   ├── reporter.py            # HTML/Markdown 报告生成
│   └── results/               # 各算子评估结果
├── optimizer/                 # 迭代优化层
│   ├── reflection.py          # 反思链 (编译日志 → 修复策略)
│   └── tuner.py               # Tiling 参数自动调优
├── skills/                    # AI 辅助开发技能库 (10+ SKILL.md)
├── docs/                      # 设计文档
│   └── design.md              # AI 应用详细设计文档
├── output/                    # 生成代码输出
└── data/vector_db/            # ChromaDB 持久化

六步 Agent 工作流

用户需求 "实现 float16 GELU"
        │
   ┌────▼─────┐   ┌──────────┐   ┌──────────┐   ┌──────────┐   ┌──────────┐   ┌──────────┐
   │ 1.语义解析│→│2.知识检索│→│3.签名设计│→│4.代码生成│→│ 5.编译验证│→│6.正确性/ │
   │(元信息提取)│  │(向量+BM25)│   │(输入输出)│   │(CoT+模板)│   │(修复≤3轮)│   │ 性能验证 │
   └──────────┘   └──────────┘   └──────────┘   └──────────┘   └──────────┘   └──────────┘
                                                                                     │
                                                                               ┌────▼──────┐
                                                                               │ 迭代优化层 │
                                                                               │ (反思+调优)│
                                                                               └──────────┘

技术方案要览

RAG 知识库

向量模型: bge-large-zh-v1.5 (中文文档) + codebert-base (代码片段)
混合检索: ChromaDB 语义检索 Top-K + BM25 关键词精确召回
知识库内容: AKG DSL 语法 · CANN API 文档 · 参考算子实现 · CCE 编译错误图谱 · Tuning Playbook

Agent 编排

框架: LangGraph StateGraph
策略: 每个步骤独立 Agent，共享 AgentState
修复策略: 编译失败 → 错误类型分类 → 匹配修复模板 → 重新编译 (最多 3 轮)

性能目标

指标	目标
Compile Pass@1	≥ 60%
Compile Pass@3	≥ 85%
Correctness Pass@1	≥ 40%
Correctness Pass@3	≥ 70%
平均性能比	≥ 70%
最优性能比	≥ 90%

许可证

本项目代码遵循 Apache 2.0 许可证。

团队

「崇理」— 队长: 蒋泽宇, 中央民族大学华为昇腾AI创新大赛银奖

PR: https://www.gitlink.org.cn/mindspore-ai/ccf-akg/pulls/1