akg-op 端到端使用指南

概述

akg-op 是 AKG_Agents 与 OpenCode 集成后的算子优化 Agent,用于生成高性能算子代码。支持两种模式:

模式 触发方式 典型场景
单算子 指定具体算子或提供代码 "帮我生成一个 relu 算子"、"优化这段 layernorm 代码"
融合 要求分析模型融合机会 "分析 Qwen2 的算子融合机会并生成融合算子"

端到端流程

┌─────────────────────────────────────────────────────────────┐
│  用户在 opencode CLI 中切换到 akg-op Agent                    │
│  (Tab 键或 /agents 命令)                                   │
└────────────────────┬────────────────────────────────────────┘
                     ▼
          ┌─────────────────────┐
          │  Phase 0  环境准备   │  检查 akg_agents 可用性
          │  & 参数确认          │  🛑 确认 framework/backend/arch/dsl
          └─────────┬───────────┘
                    ▼
        ┌───────────────────────┐
        │  是否需要融合分析?     │
        └───┬───────────────┬───┘
            │ 是            │ 否
            ▼               │
   ┌────────────────┐       │
   │  Phase 1       │       │
   │  融合分析       │       │
   │  🛑 用户选择    │       │
   │  融合机会       │       │
   └───────┬────────┘       │
           │                │
           ▼                ▼
  ┌─────────────────────────────┐
  │  Phase 2  构建任务描述        │  从代码提取算子逻辑
  │  自动验证 + 🛑 用户确认       │  → {op_name}.py
  └─────────────┬───────────────┘
                ▼
  ┌─────────────────────────────┐
  │  Phase 3  生成算子            │  选择 workflow 运行
  │  🛑 用户选择 workflow         │  (kernelgen / adaptive_search / evolve)
  └─────────────┬───────────────┘
                │
            ┌───┴───┐
            │ 失败? │──── 是 ──→ 输出失败报告,任务结束
            └───┬───┘
                │ 否
                ▼
  ┌─────────────────────────────┐
  │  Phase 4  确认结果            │  🛑 展示 generated_code.py
  │                              │  接受 / 重新生成
  │  重新生成 → 回到 Phase 3      │
  └─────────────┬───────────────┘
                │ 接受
                ▼
  ┌─────────────────────────────┐
  │  Phase 5  代码集成            │  复制生成代码到工作目录
  │  (如提供了原始代码)          │  备份 + import 替换
  └─────────────┬───────────────┘
                ▼
  ┌─────────────────────────────┐
  │  Phase 6  输出报告            │  report.md
  │  展示配置、结果、文件变更     │
  └─────────────────────────────┘

使用示例

示例 1:单算子 — 仅生成

用户> 帮我生成一个 relu 算子,输入 shape (128, 4096),fp32,在 A100 上用 triton

Phase 0: 检查环境(缓存命中,跳过安装)→ 用户确认参数 framework=torch, backend=cuda, arch=a100, dsl=triton_cuda
Phase 2: 生成 relu.py → 验证 → 用户确认
Phase 3: 用户选择 kernelgen → 运行 → 生成成功
Phase 4: 展示代码 → 用户接受
Phase 5: 复制 relu_generated.py 到工作目录
Phase 6: 输出报告

示例 2:单算子 — 优化已有代码

用户> 帮我优化 /path/to/model.py 中的 layernorm,用 Triton

Phase 0: 检查环境(缓存命中)→ 用户确认参数
Phase 2: 从 model.py 提取 layernorm → 生成 layernorm.py → 验证 → 用户确认
Phase 3: 用户选择 adaptive_search → 后台静默运行(约 15 分钟)→ 生成成功
Phase 4: 展示代码 → 用户接受
Phase 5: 备份 model.py → 添加 from layernorm_generated import ModelNew → 保存集成文件
Phase 6: 输出报告,含文件变更记录

示例 3:融合模式(首次使用,无环境缓存)

用户> 分析 Qwen2 模型的算子融合机会并实现

Phase 0: 检查环境(无缓存 → 检测/安装 akg_agents → 采集硬件/Framework/DSL → 写入缓存)
         → 用户确认参数 framework=torch, backend=ascend, arch=ascend910b4, dsl=triton_ascend
Phase 1: 分析 Qwen2 forward → 发现 3 个融合机会 → 用户选择其中 2 个
  对每个选中的融合机会:
    Phase 2: 构建融合算子任务描述 → 验证 → 用户确认
    Phase 3: 生成融合算子
    Phase 4: 展示代码 → 接受
    Phase 5: 备份 → 集成到模型代码
Phase 6: 输出汇总报告

工作目录

所有产物保存在 ~/akg_agents_logs/op_{op_name}_{时间戳}_{随机ID}/ 下:

文件 说明
{op_name}.py 任务描述(算子参考实现 + 输入定义)
{op_name}_generated.py 用户最终接受的生成算子代码
{model}_generated.py 集成后的原代码副本(含 from {op_name}_generated import ModelNew
output/{workflow}_{n}/ 每次 workflow 运行的完整输出(代码、摘要、日志)
backup/ 被替换文件的原始副本(可用于回滚)
report.md 最终报告

可选 Workflow

Workflow 策略 典型耗时 适用场景
kernelgen 迭代生成→验证→修复(SubAgent) 1-5 分钟 需求明确、快速出结果(默认
adaptive_search UCB 自适应搜索(静默模式) 10-30 分钟 希望更高质量,愿意等待
evolve 岛屿模型进化算法(静默模式) 15-60 分钟 需要多样性探索,多卡并行

生成不满意可选择换 workflow 重新生成,历史结果保留在不同子目录中互不覆盖。


架构

组件 类型 说明
akg-op Agent 主编排器:模式判定、阶段调度、用户交互
akg-env-setup Skill 环境检查、硬件/Framework/DSL 采集、依赖安装
op-task-extractor Skill 从代码或自然语言构建任务文件
vllm-ascend-operator-fusion Skill 模型融合分析
kernelgen SubAgent 迭代式代码生成+验证工作流
kernel-generator Skill DSL 感知的算子代码生成(被 kernelgen 调用)
kernel-verifier Skill 多框架、多后端正确性验证(被 kernelgen 调用)
search-workflow Skill adaptive_search / evolve 后台执行

注意事项

  • 每个阶段都有人工确认环节,不会自动执行到底
  • 生成失败时直接报告错误,不会自动尝试修复
  • 替换代码前总会先备份原文件,可从 backup/ 恢复
  • adaptive_searchevolve 以静默模式后台运行,约 1 分钟轮询一次