资讯动态

CANN算子知识库实战:让Agent先检索再生成,拒绝重复造算子

发布时间:2026/10/19 10:38:07 来源:尧图企业网站定制
“CANN知识库上线同一个算子别让Agent写第二遍”——这个标题其实挺直白的就是我们在昇腾异构计算架构CANN上做算子开发时引入AI编程助手Agent之后遇到的一个很真实的痛点Agent写算子快是快但大家让它写的算子有相当一部分是团队内部早就写过的、甚至已经在线上跑了好几个版本的东西。同一个算子被Agent在不同人手里写第二遍、第三遍质量还没有保证。所以这次把知识库真正落地核心目标就一句话让Agent在动手写算子之前先知道团队里已经有什么。这篇文章就把我们怎么做算子知识分类、怎么把可复用代码块拆出来、怎么让Agent“先检索、再生成”、以及入库验证这套流程完整分享一下适合正在做AI芯片算子开发、异构计算平台支撑或者想用Agent辅助底层软件研发的团队参考。1. 这场重复劳动是怎么产生的算子开发里的“同题异写”1.1 算子库代码量的增长远超“人肉记忆”的边界先给不熟悉CANN算子开发的朋友交代一下背景。算子就是AI计算里最基本的运算单元比如矩阵乘、卷积、归一化、各种激活函数甚至一个简单的逐元素乘加Elementwise运算都要在硬件上实现成可执行的核函数。CANN体系下的算子开发既要写Host侧的调度与Tiling计算又要写Device侧的核函数比如用Ascend C或TBE表达还要考虑数据排布、内存对齐、多核切分、流水调度这些硬约束。一个中等规模的算子库经过几代项目迭代后算子的数量会膨胀到几千甚至上万。这里头最麻烦的不是数量本身而是“相似但不相同”的变体特别多。比如同样是做规约求和输入数据的维度可能是2维、3维甚至5维数据排布可能是NCHW也可能是NHWC多核切分的策略随Shape变化而不同于是派生出一堆看起来很像、但又不能直接互相替代的实现。以前这些经验靠什么传承靠老工程师的记忆靠文档靠“我看过某某模块里有一个类似的写法”这样的口口相传。但人脑的记忆容量和检索速度是远远追不上代码库膨胀速度的。等到团队规模变大、任务并行度变高的时候重复劳动几乎是必然发生的。1.2 Agent把“写代码”变成了廉价操作却没有把“已有代码”装进来引入Agent辅助算子开发之后重复写算子这件事不但没缓解反而被加速了。原因并不难理解Agent生成代码的边际成本几乎为零你说一句“帮我写一个支持多batch的L2范数算子”它马上给你整出一个看起来功能完整的版本。于是不同小组、不同成员在面对相似需求时都可能下意识地让Agent从零生成一份而不是先去问“项目里是不是已经有类似实现了”。这里有个容易被忽略的底层问题通用大模型脑海里的“算子实现”是互联网上公开代码的平均形态并不是我们自己代码库里的最优实践。它会写一个能跑的实现但不会知道你的平台上有哪些已验证过的Tiling策略、哪些数据排布是性能陷阱、哪些边界条件曾经导致过严重事故。换句话说Agent的生成结果很“通用”但算子开发恰恰是一个极端依赖“特定平台经验”的领域。所以“同一个算子写第二遍”不只是浪费人力它还带来了更隐蔽的危害新生成的版本没有经过长期线上检验可能引入旧版本早就规避过的坑。每多写一遍团队就多承担一份回归风险。1.3 知识库到底要解决什么问题沉淀的不是代码而是“决策过程”我们最初讨论知识库方案时差点把它做成一个普通的代码仓库。那时候的想法很简单把写好的算子统一放到一个库里Agent生成前先去库里找找找不到再写。但碰了几次壁之后发现这个思路不完整。原因在于算子开发里真正值钱的不是那几百行代码本身而是代码背后的一串决策为什么Tiling块大小要选128而不是256为什么这个算子的数据排布要转换两次而不是一次到位为什么多核切分时要做动态负载均衡而不是均匀分配为什么边界条件是小于等于而不是小于如果知识库只存代码Agent即使检索到了也只能“看到”一个被封装好的黑盒无法理解决策过程。当输入数据的Shape发生变化、或者硬件规格升级时它照样不知道怎么改。所以我们最终把知识库的定位调整为沉淀“已验证的算子实现模式”和“实现模式背后的决策理由”而不仅仅是代码文件。这也是为什么我们在后续设计里特别强调分类粒度、代码块拆分以及每个入库单元必须附带性能基线和适用约束。没有这些信息检索到的代码只是个死物。2. 知识库的骨架设计分类、粒度与入库门槛2.1 算子分类按“计算行为特征”分而不是按“业务功能”分知识库第一件要做的事是把已有算子分门别类。这个分类维度的选择直接决定了Agent检索的准确率。我们一开始走过弯路试图按业务功能分比如“图像处理算子”“语音特征算子”“推荐系统算子”很快就发现这样分是错的。业务功能相同的算子在硬件上的计算模式可能完全不同而业务功能不同的算子在硬件上的计算模式反而可能高度一致。举个例子一个做图像像素值缩放的算子和一个做特征归一化的算子业务上八竿子打不着但在硬件层面它们都是典型的逐元素映射Elementwise输入输出的Shape关系、内存访问模式、多核切分方式几乎一模一样。我们后来按“计算行为特征”重新划分主体以下几类Elementwise类输入输出Shape一一对应逐元素独立计算数据局部性最好。Broadcast类输入Shape不一致需要广播语义对齐重点在于广播维度推导。Reduction类沿着某个或某几个维度做归约输出维度缩减核心难点在多核切分与跨核部分和。Gather/Scatter类输入输出维度相同但存在访存索引映射重点是索引计算与缓存命中。卷积/矩阵乘类高密度计算类涉及数据切块、流水线调度和向量化指令选择。这种分类方式的优势在于Agent面对一个新算子需求时第一步通常不是“这是什么业务”而是“这个算子在硬件上的行为模式接近哪一类”。分类对了后面的模板复用才顺理成章。2.2 最小复用单元别把整个算子当积木要拆成六个标准件知识库走近路的最大坑就是把“算子”整体当成入库和检索的单位。我们第一版就这么干结果发现Agent检索到一个相关算子后还是不知道怎么用——因为目标算子和检索到的算子之间可能只有一个维度的Shape约束不同整个拿去改又太重容易改坏。后来我们把算子拆成六个标准组件每个组件独立命名、独立入库、独立标注适用条件输入输出定义张量名称、数据类型、Shape约束、内存布局。Shape推导逻辑从输入Shape推导输出Shape的规则表达式。Tiling计算确定数据分块大小、多核映射方式、主核与从核的职责划分。核函数主循环体Device侧的核心计算代码通常是最需要保真的部分。Host侧调度逻辑任务下发、内存分配、流同步相关代码。验证与性能基线Golden Data生成方式、数值误差阈值、典型Shape下的执行时间。这六类组件分开存Agent就能按需组合——它可能复用某个算子的Tiling策略但换一套核函数主循环也可能沿用Shape推导规则但调整多核映射方式。从“复用整个算子”变成“复用组件”灵活度高了一个量级。2.3 入库条目的Schema让代码变成可检索、可判断的“结构化条目”光有分类和拆块还不够入库条目本身必须结构化。我们设计了一套Schema每个组件入库时都要填写固定字段。这个设计直接影响Agent检索命中的效果因为只有字段齐了检索时才能按多条件过滤而不是靠全文倒排硬猜。字段大致如下表字段说明示例组件类型六个标准件之一Tiling计算所属算子类别对应行为特征分类Reduction输入数据类型支持的数据类型列表fp16, fp32输入Shape范围已验证的Shape覆盖范围batch 1-64, 维度 3-5输出布局输出张量的内存排布NCHW关键约束使用时的硬性限制单核数据量不得小于256B性能基线在参考硬件上的执行时间1024x1024 fp16 平均 1.2ms已知风险点历史上踩过的坑当原维度不能被块大整除时需补零逻辑关联算子与之配套使用的组件IDKNL-1023, TIL-045这套Schema的价值在于Agent检索后不仅能拿到代码还能拿到一个“该组件是否适用于当前需求”的判断依据。输入Shape范围不匹配、输出布局不对Agent可以提前发现而不是生成之后编译报错才回头检查。2.4 入库门槛宁可短时间少入库也不要把未验证代码丢进知识库知识库刚开始运营时最容易犯的错就是“看到代码就往里塞”。团队成员都很忙入库这件事天然会被当成负担于是一旦效率优先就会出现大量质量参差不齐的条目整个库很快就废了。我们后来定了很硬性的门槛宁可入库速度慢也不放过未经完整验证的代码数值正确性必须有Golden Data对比记录数值误差在设定阈值内不允许“大概差不多”。性能达标在参考硬件上跑通关键Shape执行时间不劣于该类型算子的基线的80%。边界覆盖至少包含一个边界Shape如极端大Shape、极端小Shape、非对齐Shape的验证记录。代码规范符合团队统一风格禁止把个人实验性写法直接入库。决策说明必须附带“为什么这么写”的说明哪怕只有三五句话不能只丢一个代码块。其中“决策说明”这条最有价值也最容易被忽视。Agent生成代码时决策说明就是它的上下文没有上下文它拿到代码块还是不知道边界条件该怎么调整。我们宁可让某个组件因为说明写得太简略被打回重写也不接受一个“裸代码”入库。3. 让Agent先检索再生成接入流程与提示词模板3.1 Agent为什么能“检索到知识库”却用不好问题出在上下文注入如果只是把知识库建好但不给Agent一个明确的调用路径Agent根本不会主动去查。通用大模型的对话机制决定了它在生成代码时优先依赖的是训练阶段学到的通用知识而不是你提供的内部代码片段。想要让Agent真正“先检索再生成”必须做两件事第一在提示词里把“检索知识库”设置为一个强制步骤而不是可选步骤。第二把知识库的检索结果作为上下文片段注入到生成任务里。这一步在工程上通常走RAG检索增强生成的路线用户描述需求后先把需求文本向量化在向量库里检索最相关的组件再把检索到的组件内容拼装成上下文随提示词一起发给Agent。单纯把整个知识库文档堆进上下文是不现实的模型上下文窗口有限而且噪声太多时Agent容易抓不住重点。更合理的做法是分层检索先用向量相似度粗排找出候选组件再用Schema字段精确过滤比如输入数据类型、Shape范围必须匹配最后把命中的组件按相关性顺序组装成精简上下文。3.2 提示词模板把“检索、判断、生成、自检”一次讲清楚我们在Agent的算子开发工作流里固化了一套提示词模板。模板不能太死板但核心步骤必须稳定。实际用下来效果比较稳的版本如下你是CANN算子开发助手。任务目标是生成满足需求描述的算子实现。 严格按以下步骤执行 1. 检索知识库根据需求中的算子行为特征逐元素、归约、广播、转置、访存映射 在知识库中检索最接近的组件。必须至少找到1个参考组件并在回复中说明命中的组件ID。 2. 差异分析对比需求与命中组件的输入输出定义、Shape约束、数据布局、Tiling参数。 指出哪些可以直接复用哪些必须修改并说明理由。 3. 生成实现基于命中组件的代码模板生成新算子代码禁止从零凭空编写。复用组件时 保留原组件中经过验证的边界条件处理逻辑只在差异点上修改。 4. 自检生成完成后逐项检查输入的Shape推导、Tiling参数对齐、多核任务切分、 内存边界条件并补充至少1个边界Shape的测试用例。 5. 如果检索未命中才允许从头编写但必须说明未命中的原因并生成对应的验证计划。这套模板的核心思想是把“先检索”从口号变成流程约束。实践中Agent在步骤2的差异分析上最容易出问题——它倾向于“既然检索到了就直接生成”而跳过了“到底哪里和需求不一样”的判断。所以我们在后续迭代里会在差异分析这一步强制要求输出一个“差异对照表”用表格列出需求与命中的每个维度不允许用一句话带过。3.3 向量检索之外代码指纹检索是兜底利器纯向量检索有一个典型的盲区语义上看起来很接近的两段代码在数值实现层面可能完全不同而语义上毫无关系的两段代码在指令序列上可能异常相似。算子开发领域尤其如此比如两个不同算子的主循环体底层都是“加载数据块-计算-存储结果”的结构向量嵌入很容易认为它们高度相似但实际逻辑并不相同。为了兜住向量检索的漏检和误检我们给每个组件额外生成了代码指纹Code Fingerprint。具体来说有两种一种是AST指纹。把代码解析成抽象语法树对树结构做归一化哈希提取控制流与函数调用的模式特征。AST指纹能抓住“这段代码的结构长什么样”。另一种是Token序列指纹。把代码token化后提取关键的API调用序列、数据类型转换特征、边界条件表达式生成一个短序列标记。检索时先跑一轮向量召回再用两个指纹做交叉验证。如果向量检索召回的结果指纹不匹配就降权处理如果某个组件指纹高度匹配语义却不相关就重点人工复核。这套“向量指纹”的双通道检索比单靠向量检索的准确率明显高出一截最主要的是减少了漏检这是Agent场景最不能出现的问题。3.4 入库操作流程新算子如何一步步变成知识库里的“标准件”知识库的运营节奏是持续的小步快跑每一次验证通过的新算子实现都应该走一遍标准入库流程。我们实际执行的流程分为五步第一步自动相似度预检。新算子代码在入库前自动跑一次相似度检索看有没有和已有组件重复或高度重叠。如果相似度超过阈值不会直接拒绝入库但会提示“这可能是同一个算子的变体建议优先复用旧条目”。这一步的目的正是从源头上杜绝“同一个算子写第二遍”再次发生。第二步人工Review。由至少一名有算子库维护经验的工程师复核代码重点检查边界条件处理和性能基线标注是否真实。这一步不能省略Agent可以辅助生成但入库的最终决定权必须交给有经验的开发者。第三步补充测试数据。为组件生成标准测试用例至少包括一个常规Shape、一个极端Shape、一个非对齐Shape记录数值误差和执行时间作为性能基线的初始值。第四步结构化登记。按照前面那张Schema表格把组件的输入输出定义、Shape范围、约束、基线、风险点全部填齐。第五步发布与通知。组件入库后在团队内部广播一次附上组件ID和适用场景描述。这个看似多余的“广播”动作其实培养的是团队的先查后写习惯。久而久之开发者自己就会先搜知识库而不是直接叫Agent从零写。4. 一次真实改造复盘从“重复写”到“查一次改一次”4.1 案例背景一个L2范数算子被不同小组分别写了三次直接讲一个我们团队里真实发生的例子。项目推进过程中模块A的同学让Agent写了一个L2范数相关的归约算子功能是计算一批特征向量的模长输入Shape是二维矩阵输出是一维向量。模块B的另一个同学隔了不到两周也提出一个类似需求但输入多了一个维度当时他没有查过代码库又让Agent从零写了一个变体。真正暴露问题的是第三周模块C又要用类似算子。三个人三种写法数值结果虽然都能跑通但Tiling策略、多核切分方式、边界处理逻辑各不一样性能差异显著。最麻烦的是这三个版本里有两个的边界条件处理是有隐患的——只有在输入维度可以被块大小整除时才完全正确。这个案例其实很典型三个小组做的算子在“硬件行为模式”上是同一类但大家从Agent拿到的代码是三个不同的通用实现谁也没感觉自己和别人写的是同一个东西。4.2 改造之后的流程需求预览、组件命中、模板生成三步走知识库上线后我们把新算子开发流程改成了标准三段式。第一段Agent会根据需求描述自动生成一个“需求预览”其中包含算子类别判断、输入输出Shape轮廓、数据布局猜测。这个预览会先送到检索模块用Schema字段过滤一遍过滤的结果连同预览一起展示给开发者。第二段组件命中。比如刚才说的L2范数归约算子检索模块会先命中“Reduction类-Tiling计算组件”再命中“Reduction类-核函数主循环体组件”还会把以往测试中积累的“当原维度不能被块大整除时需补零”的风险点直接提出来。开发者看到这些命中结果基本可以确认方向对了。第三段模板生成。Agent基于命中的组件模板生成新算子同时把风险点自动写进生成的代码注释里。这里有个细节由于重建了统一的组件即使输入维度多了一维Agent也可以在原有Tiling策略上做渐进式修改而不是彻底推翻重写。实测下来一个中等复杂度的算子从需求提出到完成验证时间大约缩短了四成到五成。4.3 改造后的量化效果命中率、重复率和回归风险的全局变化我们当时给这个改造定了三个观测指标。第一个是组件命中率新算子开发过程中至少有1个组件能从知识库中直接命中的比例上线稳定后大概在七成以上。第二个是重复代码率通过采样审查新提交的算子中存在重复实现的比例这个数值稳定下降Agent重复编写的次数明显减少。第三个是线上回归风险由于新版本不再依赖未经验证的临时实现线上暴露出的边界条件类问题数量显著降低。这三个指标说明的其实是同一件事知识库的价值不在于“多了一个工具”而在于改变了算子生产的默认路径。开发者默认从检索开始Agent默认基于已验证组件生成新代码的一部分根基天然是验证过的线上风险自然被前置拦截。4.4 这个案例给团队的最大启发知识库是开发流程的一部分不是文档部门这次复盘的结论让我印象很深知识库如果只是静态的代码集合基本不会有人用只有把它嵌进AI辅助开发的完整流程作为Agent生成代码前的必经环节它才会真正生效。团队里后来形成了一个习惯需求描述里第一句话如果说是“看看知识库里有没有类似的”那这一轮的开发质量通常都很有保障。反之如果跳过检索直接让Agent生成后期review和修bug的时间往往会翻倍。5. 常见问题画像与避坑清单5.1 知识库运营中最常遇到的五个问题及排查思路知识库上线容易运营难。这半年里我们踩过的坑不少归拢下来主要是下面五类问题典型表现排查思路Agent检索命中但生成结果错误命中了相关组件生成代码在编译或运行时出错重点查提示词中差异分析是否被跳过确认检索结果是否正确注入上下文相似组件太多Agent无从选择一次召回十几个组件生成结果反而混乱调整Schema过滤条件强制加入Shape范围和数据类型过滤减少候选集规模知识库条目质量下降新入库组件缺少决策说明无法判断适用条件严格执行入库门槛缺少决策说明的组件退回补全安排周期性质量审计团队开发者不主动查库新算子开发过程完全不看知识库把检索步骤嵌入Agent提示词和开发流程不让它成为可选动作旧组件因架构升级失效硬件升级后旧Tiling策略不是最优方案对失效组件做归档而非删除保留历史版本供存量项目查询同时标注“不再推荐”第一类问题最隐蔽。Agent命中了正确的组件但如果提示词没有强制它说明“命中的组件与需求之间的差异点”它很容易把组件里的固定参数原样带到新代码里然后在新Shape下暴露边界问题。排查这类问题时先看它的生成说明里有没有做差异对照基本就能定位到根源。5.2 知识库入库的三个“不要”不要当垃圾桶、不要当万能库、不要当静态文档第一个“不要”是不要把知识库当垃圾桶。凡是测试没过、性能不达标、或者只是“先放着以后再改”的代码一律不能入库。我们处理过几轮“垃圾条目”清理发现每一条垃圾代码都会在未来某个时刻被Agent检索到然后被当座“参考实现”生成到新代码里浪费的是整个团队的时间。第二个“不要”是不要把知识库当万能库。知识库覆盖的是团队已验证过的计算模式它解决不了所有问题。如果需求特征和已有组件差异很大Agent应该明说“未命中需要从头生成”并给出验证方案。强行让Agent把不匹配的组件拼凑起来用只会得到四不像的代码。我们在提示词里特意留了“允许未命中”这个分支就是承认知识库的边界。第三个“不要”是不要把知识库当静态文档。硬件架构在升级算子的最优Tiling策略在变化性能基线也会被新的优化不断刷新。知识库必须在每次算子调优、每次架构升级后做定期刷新。我们一度把知识库当成了“建完就完事”的项目结果半年后不少条目的性能基线已经是过时数据。后来改成每季度做一次全库抽样审计重点核对性能基线和已知风险点。5.3 避坑心得如何避免知识库变成新的“文档孤岛”最后分享一条我个人的体会。知识库如果做成独立于日常研发流程的“另一个平台”几乎注定会被边缘化。真正让它活起来的做法是把知识库的检索入口嵌入到Agent交互的第一步让开发者在提需求的那一刻就看到“类似实现”的提示而不是等代码写完了再去查。一个很实用的小动作是在需求描述入口增加一个“相似检索预检环节”。也就是说开发者需求提交后Agent先返回一条“知识库中已有类似组件是否基于现有组件修改”开发者可以选择“基于现有组件生成”或“坚持从零写”。这个选择题看似微调实际上把“先查库”这个行为固化到了开发习惯里重复劳动的比例就是从这一步开始降下来的。还有一个小技巧给每个知识库组件绑定一两个“反例说明”也就是历史上因为错误使用这个组件踩过的坑。比如某个Tiling策略曾经因为块大小选择不当导致多核负载不均Agent生成代码时如果检索到该组件提示词会自动带上这段反例信息。基于反例做约束比单纯告诉Agent“你要注意边界”有效得多因为反例自带具体场景和具体数值模型更容易理解和遵循。知识库这件事说到底不是技术问题而是流程与习惯问题。技术方案无非是向量检索、Schema设计、代码指纹那一套但真正让“同一个算子不写第二遍”成立的是团队对“先查后写”的共识。现在Agent写代码的速度越来越快知识库的价值也会越来越大谁先把已验证的经验沉淀好谁就能在后续的开发里省下最多的时间。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑