融合算子识别与替换

这个 Skill 做什么

针对昇腾 NPU 上跑的 PyTorch 模型（训练或推理），找到可以用 torch_npu 融合算子替换的代码段，给出替换示例，再做性能和精度验证，最终给出"该不该换"的结论。

整个工作分两个阶段，阶段一输出识别报告后必须停下来等用户确认，确认后再进阶段二逐个验证。

核心概念

• 融合算子：把多个连续计算层（如 matmul→激活→归一化）合并成一个 NPU 硬件内核，中间结果不写回显存。典型：npu_fusion_attention 把 Q@K^T→mask→softmax→@V 合成一个算子。
• 亲和算子：不改数学逻辑，仅把 NPU 不擅长的随机访存指令等价代换成连续乘加。典型：IndexPut→乘法、where→lerp。
• 亲和 API：NPU 官方高阶封装函数，内部串联多个亲和算子。典型：NpuFusedAdamW。

本 Skill 主要处理融合算子替换；亲和算子/亲和 API 在"高频算子序列"环节顺带识别。

全局代码替换规范（严格执行）

无论是阶段一报告中的代码示例，还是阶段二实际修改的代码，所有生成的替换代码必须满足以下生产级条件，绝不允许使用伪代码：

1. 非破坏性修改 (Fallback 机制)：必须通过参数（如 args 或 model.config）或 is_torch_npu_available() 进行条件判断，必须原样保留原代码作为 else 分支。
2. Dtype 严格对齐与 AMP 意识：必须保证融合算子分支的输出 dtype 与原生分支完全一致。注意记录并保持原有的 Autocast 状态，防止因替换 API 导致异常的类型强转（注意 float16/bfloat16/float32 在特定 NPU 算子下的支持度差异）。
3. Layout 与连续性保障：NPU 融合算子对 Tensor 内存布局敏感。在传入融合 API 之前，需确保输入 Tensor 是连续的，适时使用 .contiguous()，并在 Fallback 分支中考虑其影响。
4. 优先复用：优先复用目标仓库内已有的开关参数、全局环境变量或 wrapper 命名。

两阶段工作流

阶段一（识别，必做）                      阶段二（逐个验证，按用户确认的清单）
┌──────────────────────────┐    用户确认   ┌──────────────────────────┐
│ 1. 常用融合算子扫描（必做）│ ──────────► │ 1. 按条件分支改代码        │
│ 2. 高频算子序列识别（advanced）│  暂停     │ 2. 逐个跑：精度优先        │
│ 3. 融合价值评估            │  不改代码   │ 3. 精度不过→直接否决        │
│ 4. 输出 markdown 识别报告  │            │ 4. 精度过→再看性能         │
└──────────────────────────┘            └──────────────────────────┘

---

阶段一：识别可替换的融合算子

1.1 常用融合算子（必做项）

拿着模型代码，对照下表，看有没有匹配得上的代码逻辑：

手写模式	融合 API	一句话要点
手写 RMSNorm（x.pow(2).mean → rsqrt → * weight）	torch_npu.npu_rms_norm(x, weight, epsilon=eps)	返回 tuple，取 [0]；最简单、收益稳定
a1 * F.silu(a2) 门控乘法	torch_npu.npu_swiglu(x, dim=-1)	concat 顺序必须正确：[a2, a1] 让 SiLU(a2)*a1
手写旋转位置编码（t*cos + rotate_half(t)*sin）	torch_npu.npu_rotary_mul(t, cos, sin)	cos/sin 需 expand_as(t_)；处理 t_pass_ concat 回来
手写 Q@K^T → mask → softmax → @V	torch_npu.npu_fusion_attention(...)	必须显式 causal mask；atten_mask 中 True=屏蔽；返回 tuple 取 [0]
MoE expert 循环 matmul	npu_grouped_matmul / Ops.gmm + npu_moe_token_permute/unpermute	需确认 expert 归属、weight layout、空 expert、routing weight

每个算子的详细识别 pattern、原始→NPU 完整代码对照、踩坑点和 concat 顺序推导，见 references/common_fusion_ops.md。 在写替换代码前务必读对应小节，并用 scripts/query_torch_npu_api.py 查询当前环境的 API 说明文档确认签名、参数约束和返回值；脚本查不到时再看文档链接或在线文档兜底。同时在目标仓库内搜索已有融合算子实现，优先参考项目内已验证的 wrapper、开关命名、layout/mask 处理和 fallback 写法。

1.2 高频算子序列识别（advanced，基于 profiling）

如果用户提供了替换前的 profiling 数据（或主动采集了），基于 profiling 里的高频算子序列，能发现常用清单覆盖不到的融合机会。在 kernel_details.csv 里找属于同一模块、上下衔接的算子序列（如 slice+matmul+gelu），按 pattern 去知识图谱或 torch_npu API 列表找现成融合 API。注意算子序列必须属于同一模块、上下衔接，不要跨模块拼凑。

如何从 profiling 判断这段代码到底是不是融合有价值的瓶颈，见 references/fusion_value_assessment.md（含 Roofline 加速上限表、Memory/Compute/Launch-bound 现象对照）。

1.3 限制条件与边缘场景排查（重点防范）

在识别出候选算子后，必须评估以下边缘场景：

• 动态 Shape 敏感性：评估该代码段是否涉及动态 Shape（如每次输入的 Sequence Length 均不一致）。若存在动态 Shape，NPU 算子可能会发生频繁 Re-compile 严重拖慢性能，需在报告中明确预警并评估是否值得替换。
• 连续性要求：排查输入 Tensor 是否经过切片（Slice）、转置（Transpose）等破坏内存连续性的操作

1.4 输出格式

阶段一只输出 Markdown 识别报告，不修改模型代码。报告必须包含：汇总表、逐个候选分析（、建议顺序，以及符合全局代码规范的生产级代码替换示例。具体参考 references/identification_report_format.md。

1.5 阶段一边界

阶段一结束后不要直接进阶段二改代码。先输出 Markdown 识别报告，让用户确认认可哪些替换、做哪些。用户确认后再进阶段二。

---

阶段二：修改代码 + 逐个验证

进入阶段二的前置条件：用户输入必须包含可以实际运行模型的脚本、命令或等价拉起方式。如果没有，先向用户索要。

2.1 代码修改与适配

严格遵循前文定义的【全局代码替换规范】修改代码。示例：

args = get_args()
if args.use_fused_rmsnorm and is_torch_npu_available():
    # aclnnRmsNorm 支持 FP32 gamma，对 FP32/BF16 输入都成立
    output = torch_npu.npu_rms_norm(hidden_states, self.weight.float(), epsilon=self.variance_epsilon)[0]
else:
    # 原 fallback 分支原样保留
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
    output = (self.weight * hidden_states).to(input_dtype)
return output

控制方式优先级：项目已有参数 > 新增参数 > 环境变量 > 仅判断 is_torch_npu_available()。

2.2 评测原则：精度优先，逐个对比 baseline

第一步：算子级精度对齐 (Unit-level Validation) 在跑端到端训练前，先造一组与模型实际 input shape/dtype 完全一致的随机 Tensor，分别送入原生算子和 NPU 融合算子，对比输出结果的差异。

• 标准：Cosine Similarity > 0.999 或 Max Absolute Error (MAE) < 1e-4。
• 若单算子级未对齐，直接判定为不通过，无需进入后续 E2E 测试。

第二步：端到端精度验证 (E2E Loss Validation) 算子对齐后，跑 50+ 个迭代，要求修改前后 Loss diff 在 1% 以内。如果早期迭代就有明显精度溃散，提前结束。精度不满足 → 直接标注不通过，跳过性能评测。

第三步：性能验证 精度通过后测性能，默认不开 profiling。主要关注端到端 step_time：是否有稳定下降。

2.3 数据记录与报告生成

验证期间从日志提取逐 step 数据（step id / loss / step_time）写入 JSONL；每个候选验证完后复制 assets/validation_report_template.html 模板、替换 REPORT_DATA 生成 HTML 报告。详细流程、JSONL 格式、REPORT_DATA 字段说明见 references/verification_workflow.md。

2.4 长任务过程记录

用 TaskList 跟踪每个候选替换的验证状态（待测 / 精度验证中 / 精度不通过 / 性能验证中 / 已结论），结论写进本地 markdown，方便中断后续接。

---

相关脚本

• scripts/query_torch_npu_api.py — 查 API 签名/约束，写替换代码前必用
• scripts/merge_jsonl.py — 合并 baseline+fused JSONL 为模板 metrics 数组
• assets/validation_report_template.html — 验证报告 HTML 模板，复制后替换 REPORT_DATA
• 算子信息在线文档：https://gitcode.com/Ascend/op-plugin/tree/master/docs/zh/custom_APIs/torch_npu

---

决策清单

• 用户给的是模型代码/拉起脚本？有没有替换前 profiling 数据（有则进 1.2）？
• 常用融合算子清单逐个扫了吗？仓库内已有融合实现搜过了吗？
• 替换 API 的签名/约束查了 query_torch_npu_api.py 而非凭记忆吗？
• 阶段一报告输出后，停下来等用户确认了吗（没直接改代码、没生成 HTML）？
• 阶段二逐个验证，精度优先，精度不过直接否决了吗？
• 逐 step 数据写进 JSONL 了？验证完用模板生成 HTML 报告了？
• 长任务过程记录和 TaskList 跟上了吗？