SuperAPI AI 工具集

external-cannbot-model-model-infer-fusion

编程与开发

查看项目

更新于 10-15 17:19

技能介绍

用途:分析 PyTorch 模型代码中的计算模式,匹配 torch_npu 融合算子进行替换优化。依托 torch_npu 本地 docstring 查询脚本 scripts/torch_npu_query.py 与仓库已有模型的融合算子使用经验,按模块逐一分析、匹配、替换、验证。

主要能力与流程:

  1. 分析模型代码并拆解模块:识别顶层结构(Embedding→Transformer Blocks→LM Head),拆解 Attention、MoE/FFN 等模块,记录 Prefill/Decode 差异、量化需求(BF16/W8A8/W8A8C8/W4A16)、分布式配置(TP/EP/DP)及 layout、cache 格式、metadata 等约束,展开到可替换链路级别。
  2. 按模块匹配仓库参考实现:Attention 区分 GQA/MHA 与 MLA(含 Indexer),结合 cache 形态(连续非PA/PA)、写入语义与 Attention Core 选择参考链路;MoE 区分 Dense FFN 与 MoE,按 gate 形式与并行/阶段匹配路由与专家计算链路。
  3. 查阅算子接口文档:优先本地 docstring,用以下命令查询,确认可用性与适配性: python3 scripts/torch_npu_query.py show <api_name> python3 scripts/torch_npu_query.py search "<keyword>" python3 scripts/torch_npu_query.py list [--prefix npu_] 检查 shape、dtype、layout、cache 组织及 metadata;可通过前置改造解决时标记为候选并说明改造项,硬约束不满足才标为不适配并注明依据。
  4. 分析阶段审查:模块拆解完整、参考匹配完整、算子约束已确认、候选清单已形成;仅要求分析时到此输出结果与验证计划。
  5. 逐模块实施替换:每次只改一处,验证通过再继续;做精度对比与性能对比,失败则回退并重新分析或记录新需求与阻塞原因。

适用场景:torch_npu 融合算子替换、MoE/Attention/FFN/Norm 等模块的推理算子适配。

限制:具体支持算子范围、版本与性能收益以实际环境为准;资料未提供收费、价格与平台信息。参考文档包括 references/ 下 Attention、MoE、RotaryEmbedding 等参考链路,以及 op-plugin 仓库 gitcode.com/Ascend/op-plugin。