CANN/cannbot-skills模型推理融合算子优化

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsname: model-infer-fusion description: 基于 PyTorch 框架的昇腾 NPU 模型推理融合算子优化技能。分析模型代码识别可替换为 torch_npu 融合算子的计算模式生成替换方案。触发场景torch_npu 融合算子替换、MoE/Attention/FFN/Norm 等模块的推理算子适配、torch_npu API 使用咨询。基于仓库已有模型的融合算子经验按计算语义推荐最佳方案。torch_npu 融合算子优化技能分析 PyTorch 模型代码中的计算模式匹配 torch_npu 融合算子进行替换优化。基于项目内收录的 torch_npu 官方算子接口文档references/torch_npu_API/和仓库已有模型的融合算子使用经验按模块逐一分析、匹配、替换、验证。工作流程第一步分析模型代码拆解模块分析模型代码输出模块清单。拆解网络结构识别顶层结构Embedding → Transformer Blocks → LM Head拆解每个 Block 为独立模块Attention 层Norm → QKV 投影 → RoPE → KV Cache → Flash Attention → O 投影MLA 含 V absorbMoE / FFN 层Norm → Gate → 路由分发 → 专家计算 → 聚合其他模块Embedding、LM Head、跨层残差等记录运行场景特征Prefill / Decode 分支差异同一模块在两个阶段可能走不同算子路径替换时分别处理量化需求BF16 / W8A8 / W8A8C8 / W4A16分布式配置TP / EP / DP影响 MoE 路由算子选择固定约束layoutNZ/BSND/TND、cache 格式、metadata关键子链路展开对 Attention、MoE 等复杂模块展开到可替换链路级别AttentionRoPE、KV Cache 写入/读取、Attention CoreFA / FA v2 / Sparse FAMoEGate / TopK、Routing Init / Dispatch、Expert 计算、Finalize / Combine、Shared Expert若有其他关键链路Residual Norm、QKV Projection、O Projection、MC2 / AllToAll 等一并纳入第二步按模块独立匹配仓库参考实现对第一步拆解的每个模块独立在仓库参考实现中匹配最接近的算子链路。命中后将该路径作为候选蓝图并对照第一步的关键子链路清单补充分析该路径未覆盖的部分未命中的模块跳到第三步在算子总表中搜索。Attention 层注意若 FA 融合算子已在前置阶段替换仅跳过 FA 调用本身。Attention 的其他子链路RoPE、KV Cache 写入、ResidualNorm、QKV/O Projection 等仍需逐一检查是否有可用融合算子。同时注意多步骤整体融合算子如npu_kv_rmsnorm_rope_cache融合 RMSNormRoPECache写入npu_mla_prolog_v3融合 Q/KV投影RMSNormRoPECache写入不要将单个子模块简单归为标准线性层无融合算子而跳过评估。分析 Attention 时KV Cache需要与Attention Core结合讨论不建议完全拆开除架构外还应同时确认cache 组织方式连续非PA / PA写入索引语义kv_len / start_pos或slot_mapping / cache_index后续Attention Core非融合 attention、npu_fused_infer_attention_score、npu_fused_infer_attention_score_v2、npu_sparse_flash_attention常见实现组合包括连续非PA 连续写入如scatter_update_kv_len / start_pos→ 再评估非融合 attention 或融合FAPA slot-based write如npu_scatter_nd_update_/npu_scatter_pa_kv_cache/cache_index→ 再评估与该 cache 形态匹配的FAMLA absorb 特化路径npu_mla_prolog_v3(cache_index / slot_mapping) 融合FA先根据当前实现的架构、cache 形态、写入方式和Attention Core组织判断它更接近哪条参考链路再进入对应详情文档。Attention 架构 │ ├─ GQA / MHA标准多头 / 分组查询注意力 │ │ │ ├─ Prefill: 当前 batch 的 q/k/v → FA(sparse_mode3 推荐sparse_mode2 不推荐) │ └─ Decode: 从 KV Cache 读取 → FA(actual_seq_lengths_kv) │ → 详情references/module-attention-gqa.md │ └─ MLAMulti-head Latent Attention低秩 KV 压缩 │ ├─ 无 Indexer │ ├─ Prefill: 分步投影 → 展开 K/V → FA v1 或 v2 │ └─ Decode: absorb手动或 npu_mla_prolog_v3→ FA v1 或 v2 → V absorb │ → 详情references/module-attention-mla-absorb.md │ └─ 有 Indexer稀疏 Top-K KV 选择 ├─ Prefill/Decode 共路径 │ npu_mla_prolog_v3 → Indexer → 稀疏 FA → V absorb → 详情references/module-attention-mla-indexer.mdMoE / FFN 层先确认当前模块属于 Dense FFN 还是 MoE若为 MoE再结合 gate 形式、并行模式和阶段判断更接近哪条参考链路。是否有 MoE │ ├─ 无Dense FFN→ 在算子总表中确认 Dense Linear / Activation / Norm 可用融合算子 │ └─ 有 MoE │ ├─ Gate 算子 │ ├─ softmax 打分 → npu_moe_gating_top_k_softmaxqwen3-moe │ └─ sigmoid/noaux → npu_moe_gating_top_kdeepseek 系列 │ └─ 路由 专家计算按并行模式和阶段区分 ├─ Prefill纯 TP: init_routing_v2 → grouped_matmul → finalize_routing ├─ PrefillEP: init_routing_v2 → AllToAll → re_routing → grouped_matmul → finalize_routing └─ DecodeEPTP: MC2 dispatch_v2 → grouped_matmul → MC2 combine_v2 → 详情references/moe-fusion-reference.md未匹配模块其他未被上述判断树覆盖的模块Embedding、LM Head、跨层残差、Diffusion 特有模块等跳到第三步在算子总表中搜索。第三步查阅算子接口文档确认可用性与适配性对第二步命中的算子或未匹配到模式的模块查阅 torch_npu 官方算子接口文档算子总表references/torch_npu_API/torch_npu_list.md本地索引算子详情op-plugin 在线文档 — 参数说明、dtype/shape 约束、代码示例确认可用性对模式命中的算子查阅详情文档确认参数约束和适配场景对未命中模式的模块在算子总表中搜索阅读详情文档分析功能适配验证检查 shape、dtype、layout、cache 组织及 metadata 是否满足算子要求若差异可通过合理前置改造解决如格式转换、RoPE 预计算与取值路径整理、KV Cache 静态化/PA 改造等应标记为“候选 需前置改造”并说明所需改造项部分流程较复杂的改造可询问用户是否采用仅当差异属于硬约束且无法通过合理前置改造解决时才可标记为”不适配”。标记时必须注明具体硬约束如算子报错信息、文档明确的参数限制不能仅以”需改动较大”为由标记不适配第四步分析阶段审查在进入代码替换前审查以下各项。未完成的项须返回对应步骤补齐不得跳过直接进入实施。若当前任务仅要求分析则在本步结束输出分析结果、候选方案和验证计划不进入代码实施。分析阶段审查项模块拆解完整已展开到可替换链路级别含 Prefill/Decode 分支差异Attention/MoE 关键子链路已覆盖参考匹配完整已按模块匹配仓库参考实现或确认无匹配已补充分析参考路径未覆盖的子链路算子约束已确认已查阅 reference / API 文档确认候选算子的适配约束和量化兼容性候选清单已形成每个候选模块明确前置条件、最小验证切口和阻塞点第五步逐模块实施替换已有全面的算子候选分析后依照替换流程对候选清单中的所有模块逐个迭代替换每次只改一处验证通过再继续下一个不得跳过任何已进入候选清单的模块。若当前模块无法继续实施也必须记录其失败证据、阻塞原因和当前结论。若候选模块依赖以下前置改造且尚未完成可按需查阅对应资源RoPE预计算、缓存和取值路径整理references/rotary-embedding-pattern.mdKV Cache静态化 /PA改造参考model-infer-kvcacheskill融合算子替换流程完成前置改造如需要替换该模块的算子代码验证精度对比融合前后输出对齐验证性能对比确认有收益若验证通过 → 保留继续下一个模块若验证失败 → 回退改动重新分析该模块是否有其他可用算子或替代方案→ 尝试替代方案确认当前无适配算子但有融合收益 → 记录为新需求模块位置、计算语义、不适配原因、期望融合形式当前无法继续实施 → 回退后记录失败证据和阻塞原因继续下一个记录该模块的优化报告精度对比结果、性能对比结果、日志或报错路径注替换时可参考仓库中最接近的模型实现和references/下的算子接口文档使用说明。实施阶段检查项候选清单中的所有模块均已逐一处理替换或记录跳过原因无遗漏每个已实施模块均记录了精度和性能对比结果实施失败的模块已记录报错信息、尝试过的处理方式和最终结论参考文档索引以下文档按需查阅避免一次性加载消耗 token主题路径torch_npu 算子总表本地索引references/torch_npu_API/torch_npu_list.mdtorch_npu 算子详情在线op-plugin docsAttention: GQA 参考链路references/module-attention-gqa.mdAttention: MLA Absorb 参考链路references/module-attention-mla-absorb.mdAttention: MLAIndexer 参考链路references/module-attention-mla-indexer.mdMoE 算子模式详解references/moe-fusion-reference.mdRotaryEmbedding 预计算与调用模式references/rotary-embedding-pattern.mdop-plugin 仓库在线gitcode.com/Ascend/op-plugin【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考