在深度学习模型高效规模化发展趋势下，融合算子是昇腾生态的核心优化技术，能整合多小算子计算逻辑，降低内存开销与内核启动成本。PTA作为连接PyTorch与昇腾硬件的关键桥梁，其算子适配与上仓规范直接影响技术落地。
下面我将结合工程实践，从融合算子原理出发，拆解Permute & Unpermute算子机制，梳理Op-Plugin仓结构化适配流程，呈现从需求评审到主仓合入的全链路规范。内容涵盖关键技术点与代码、图表示例，助力新老开发者快速上手，高效完成融合算子适配与上仓，保障PyTorch模型在昇腾NPU稳定高效运行。
一、融合算子原理和使用
（一）融合算子原理
融合算子是性能优化核心技术，核心是将多个小算子的计算逻辑整合到单个内核函数中。传统小算子需单独调用内核、传递中间结果，存在显存占用高、读写频繁、内核启动开销大等问题。
融合后可减少中间结果存储与内存读写，降低带宽压力，避免数据冗余搬运，提升计算密度。
（二）融合算子使用
目前PTA支持的融合算子主要有这几个，适配的版本和数据类型见下：
另外MindSpeed里的一些算子，PTA也在同步支持：
MindSpeed已经有这些算子了，为什么PTA还要同步支持？
主要是两个原因：
- 一是能原生支持PyTorch+Torch NPU，不用额外装其他依赖，环境配置更简单；
- 二是针对FSDP场景，客户直接调用PTA接口就能快速部署，不用再做额外适配。
二、Permute & Unpermute算子介绍
（一）重排与反重排的原理
这两个算子主要用在MoE（混合专家模型）里，核心是解决多专家并行计算的数据调度问题。
重排（Permute）
为了高效地并行计算，通常会将属于同一个专家的样本聚集在一起，这种操作就涉及到对原始输入数据的重排
# 按专家编号排序，得到分组后的索引 sorted_indices = torch.argsort(indices) # 重排输入数据，相同专家的样本连续 sorted_inputs = inputs[sorted_indices] # 统计每个专家处理的样本数 counts = torch.bincount(indices, minlength=num_experts) # 按专家拆分数据 split_inputs = torch.split(sorted_inputs, counts.tolist())反重排（Unpermute）
在专家处理完各自的样本后，将输出结果按原始输入顺序重新组合，以便后续网络层能正确接收
# 拼接所有专家的输出结果 outputs = torch.cat([experts[i](split_inputs[i]) for i in range(num_experts)]) # 恢复原始输入顺序 final_outputs = outputs[reverse_indices]整个流程其实就是：输入数据经过门控网络分配专家编号，先Permute重排让同专家样本聚集，专家并行计算后，再通过Unpermute恢复原始顺序，最后输出结果。
（二）aclnn算子接口
目前Op-plugin算子仓库主要支持两类算子：ACLOP和ACLNN，它们的定位和用法差别见下：
因为PTA适配大多是结构化适配，也就是直接调用这些已有的aclnn接口，所以得把接口逻辑摸清楚。
主要从这几个维度去了解：
- 算子功能（知道它能干嘛、用在什么场景）
- 计算公式（搞懂底层实现原理）
- 参数说明及返回值（处理好接入时的上下文）
- 约束说明（清楚算子的使用限制）
- 调用代码（看C++源码，梳理整体逻辑）。
（三）算子引申：Routing Map版本
现在我们适配的Token Permute是基于Megatron 0.8.0版本开发的，而Megatron 0.12.0版本里出了个Routing Map版本的Token Permute，叫MoeTokenPermuteWithRoutingMap，相比于原来的版本比有不少改进：
原来的版本里，Permute和Unpermute都要基于expert_indices重新计算映射关系，有点重复劳动；而Routing Map版本只需要计算一次routing_map，Permute和Unpermute都能复用，不仅灵活度高了，性能也可能更好。
三、Op-plugin仓介绍&结构化适配
（一）Op-Plugin仓介绍
Op-Plugin其实是PTA的算子插件，以第三方的形式提供NPU算子库调用能力。它是基于昇腾的Ascend Extension for PyTorch（torch_npu）开发的，所以编译和使用之前，得先了解并安装好昇腾的torch_npu，具体可以参考昇腾社区的使用手册。
在算子编译调试的时候，OpPlugin库支持单独编译，能对应不同的PTA和Python版本。比如要编译v2.1版本，步骤大概是这样：
# 拉取对应分支的代码 git clone -b master https://gitee.com/ascend/op-plugin.git # 进入插件根目录 cd op-plugin # 执行编译构建，这里以Python3.8和PyTorch v2.1.0为例 bash ci/build.sh --python=3.8 --pytorch=v2.1.0目前它支持torch_npu 2.1、2.3、2.4、2.5、2.6、2.7、2.8以及master版本。
仓库结构
因为Permute和Unpermute已经有对应的aclnn算子接口了，所以我们用结构化适配的方式就行——编写好算子接口调用逻辑、前反向绑定规则和shape校验规则，就能在opapi目录下自动生成cpp算子文件，不用手动写完整的实现代码。
（二）结构化适配
1. 适配流程
2. 算子逻辑适配
以Permute算子为例，前向适配的配置大概是这样：
Unpermute算子的适配逻辑和Permute类似，都是先定义好函数、输出参数，计算中间参数，最后调用对应的aclnn接口执行，见下：
3. 输入输出shape检查
shape检查是很重要的一步，能提前规避很多参数错误。比如Permute前向的shape检查：
Unpermute和梯度算子的shape检查逻辑也类似，核心就是校验输入张量的维度是否符合要求，再根据输入参数计算出输出的shape。
4. 前反向绑定
前反向绑定的作用是让自动微分系统（Autograd）能正确构建反向传播图，比如Permute和Unpermute的绑定配置：
简单说就是告诉系统，计算某个前向算子的梯度时，应该调用哪个反向算子。
5. UT测试用例
UT测试主要是验证算子功能和精度，Permute和Unpermute的测试主要包括三部分：
- 小算子实现：手动用PyTorch原生算子实现相同功能，作为比对的基准；
- 单算子比对：构建不同的测试用例，分别跑小算子和融合算子，比对两者的结果，正反向都要测，反向用.backward()验证；
- 整体功能验证：因为Permute和Unpermute是镜像操作，所以把输入经过“Permute→Unpermute”后，应该和原始输入一致，用这个来验证端到端功能。
后续如果遇到算子精度问题，也能通过这些UT测试用例排查，看看是参数传递、shape计算还是接口调用出了问题。
四、上仓流程与PTA主仓合入
（一）整体流程
（二）CLA签署
提交代码贡献之前，还有个重要步骤是签署CLA（贡献者许可协议），这是开源项目的常规要求，主要是明确贡献者和项目方的权利义务，保障知识产权合规。
CLA分两种：个人CLA适合个人开发者，自己签署就行；企业CLA是公司签署的，能覆盖公司所有员工的贡献，不用每个人单独签。
比如之前提交代码的时候，会有自动化工具检查是否签署了CLA，如果没签会提示，签完之后工具会自动更新状态，代码才能进入后续评审流程。
（三）CI测试
CI测试是保障代码质量的关键，提交代码后会自动触发一系列测试：
- Code Check：检查代码规范性，比如命名、格式、注释这些是否符合项目要求；
- Build：基于不同版本的PyTorch、torch_npu，还有x86、ARM两种架构做编译测试，确保兼容性；
- UT：在不同版本上跑所有单元测试，验证功能和精度。
测试的版本覆盖很全，包括master、v1110、v210、v251、v260、v271这些，只有所有测试任务都成功了，才能进入下一步。
（四）PTA主仓合入
提供算子设计文档和完成meta注册后，会排入下一个商发版本 在商法版本正式发布前，需要使用算子接口有两个方式：
- 通过CMC获取对应版本的每日构建/转测包；
- 基于现有的商发版本PTA，手动补充对应算子逻辑，编译安装
- 拉取对应版本的PTA和Op-Plugin代码
- 合入对应算子逻辑（前反向绑定，算子逻辑，shape检查）
- 编译PTA包并安装测试
总结
本次 PTA 融合算子结构化适配与上仓实践收获颇丰：
- 理清了融合算子原理、Permute/Unpermute 工作机制及 ACLOP 与 ACLNN 算子差异；
- 熟悉了从需求评审、开发适配到测试合入的完整规范流程，可复用经验；
- 掌握 UT 测试用例设计与精度问题定位方法，能结合客户模型识别算子需求；
- 明确性能优化思路：优先替换融合算子，再通过 profiling 等分析高耗时算子针对性优化。
最后，整个知识体系其实可以归纳为四大块：融合算子基础、核心算子解析、结构化适配、上仓流程。把这些内容掌握好，就能高效完成融合算子的适配与上仓，让PyTorch模型在昇腾NPU上跑得更高效。
注明：昇腾PAE案例库对本文写作亦有帮助。
)
