前言CANNCompute Architecture for Neural Networks是昇腾NPU的底层计算框架承担着将深度学习算子高效映射到硬件的职责。在Transformer架构主导大模型时代的背景下标准算子组合方式在昇腾NPU上面临严重的访存瓶颈——注意力计算中QKV投影、缩放、Softmax、输出投影之间的中间结果频繁写入全局存储每个算子独立调度又引入大量核函数启动开销。ops-transformer算子库正是为解决这类问题而生定位于CANN大模型推理与训练场景的专用算子集计划于2025年11月全面上线。仓库以子目录形式组织了attention、ffn、gmm、moe、mhc、mc2六大类算子每类针对Transformer模型中不同模块的访存特征和计算模式进行优化。attention子目录处理自注意力与交叉注意力的融合计算ffn子目录覆盖前馈网络的逐层融合gmm和moe子目录面向混合专家架构的分组矩阵乘法与路由调度mhc子目录涉及多头上下文相关的计算加速mc2子目录解决多通道数据复制在并行生成中的带宽瓶颈。这六类算子的共同设计目标是降低Transformer模型在昇腾NPU上的访存开销、提高片上数据复用率从而在有限的片上存储容量内压榨更高的计算密度。仓库通过gitmodules引入catlass模板库将算子内核实现与底层矩阵乘法模板解耦使得每种融合算子可以专注业务逻辑而将高性能GEMM/IMAGEMV等基础计算委托给模板层完成。attention子目录结构与Softmax融合算子演进Transformer注意力机制的计算流程可以拆解为线性投影、缩放点积、Softmax归一化、加权求和四个阶段。在标准实现中这四个阶段各自对应独立的算子调用中间张量在各级存储之间来回搬运。CANN的ops-transformer算子库将这一流程封装为融合算子以SMLANSoftmax-Linear-Attention-Normalized和SMLASoftmax-Masked-Linear-Attention两个系列为代表形成了多版本演进的格局。SMLAN系列适用于无mask约束的标准注意力计算场景将QKV投影到Softmax输出的全部流程融合到单个kernel内部SMLA系列在SMLAN的基础上增加了mask处理分支用于自回归解码阶段的上三角mask注入。两个系列均按照head维度进行拆分每个融合kernel负责处理一组attention head的完整计算以此降低kernel启动次数并提高片上L1缓存的数据局部性。mask机制在融合算子中的处理方式直接影响分支效率和数值稳定性。SMLA系列在Softmax归一化之前执行mask操作将无效位置设置为负无穷大值使得Softmax归一化后这些位置的权重趋近于零。在融合kernel内部mask通过一个布尔张量传入kernel在缩放点积完成后检查每个位置的mask标志位对被mask的位置写入极小值。这种处理方式避免了额外的mask乘法kernel调用但引入了分支分歧——同一warp内的线程可能因为mask分布不均匀而产生指令分叉。为了缓解这一问题SMLA系列的部分版本采用两阶段策略阶段一对所有位置执行完整的缩放点积和Softmax计算阶段二根据mask标志位将对应位置的输出清零。这种策略消除了Softmax阶段的分支分歧代价是少量无效计算。padding逻辑同样对融合算子的计算效率产生实质影响。在变长序列的场景中短序列会被填充到固定长度padding位置参与注意力计算但最终结果被丢弃。如果不加处理融合kernel会对padding位置执行完整的QKV投影、缩放点积和Softmax计算浪费算力和带宽。ops-transformer的attention算子通过传入一个有效长度参数来标记每个序列的实际token数量kernel内部依据该参数跳过padding位置的点积计算直接在输出中保留零值。这种处理方式要求序列长度参数在编译期或运行期可被kernel获取因此在不同padding策略下需要切换对应的kernel变体。// SMLA融合kernel中mask与padding的简化处理逻辑for(inti0;iseq_len;i){floatsdot(q[h],k[h][i])*inv_sqrt_d;s(ivalid_len)?-1e9f:s;// padding位置置负无穷smask[h][i]?s:-1e9f;// mask位置置负无穷attn[h][i]expf(s);}floatsum0.0f;for(inti0;iseq_len;i)sumattn[h][i];for(inti0;iseq_len;i)attn[h][i]/sum;floato0.0f;for(inti0;iseq_len;i)oattn[h][i]*v[h][i];GMM与MC2新型算子的设计动机与协同关系混合专家架构Mixture of Experts将Transformer的前馈网络替换为多个专家子网络通过门控路由机制为每个token选择若干专家执行计算。这种架构在扩展模型容量的同时保持了推理时的计算成本可控但对算子实现提出了新的要求门控路由的稀疏性导致不同token激活值在不同专家之间分配不均匀标准的矩阵乘法算子无法高效处理这种分组访问模式。GMMGroup Matrix Multiply算子正是为解决这一矛盾而设计它在矩阵乘法内核中引入分组维度将输入激活值按照门控路由结果分组到对应的专家权重矩阵上执行乘法。GMM的核心思路是将稀疏的路由决策转化为稠密的分组矩阵乘法——路由器输出的专家索引被转化为分组标识kernel依据分组标识将激活值和权重加载到对应的计算单元上每组内部的矩阵乘法遵循标准的分块乘法策略。MC2Multi-Channel Copy算子的设计背景是自回归解码阶段多token并行生成时的激活值复制问题。在投机解码或推测采样等场景中模型需要同时验证多个候选token每个候选token共享相同的KV缓存和部分中间激活值。标准的数据复制需要多次从全局存储读取相同的源数据并写入不同的目标地址在多token并行度较高时会产生大量的冗余访存。MC2算子将多通道复制操作融合为一次读取、多次广播的kernel片上存储中只需加载一次源数据即可向多个目标通道分发减少了重复的显存读取次数。GMM与MC2在组合使用时存在明确的依赖关系与调度约束。在MoE推理中GMM完成分组矩阵乘法后输出各专家的计算结果这些结果需要根据路由信息回写到对应token的输出缓冲区。当采用多token并行验证时MC2负责将共享的激活值复制到各token的计算通道GMM则在各通道内分别执行分组乘法。调度约束体现在执行顺序上MC2必须在GMM之前执行确保各通道的输入数据就绪GMM的分组标识需要与MC2的通道编号对齐避免跨通道的数据混叠。在昇腾NPU的任务流调度层面MC2和GMM被编排为同一Stream内的先后两个kernel中间无需Host侧的同步干预由Device侧的事件依赖机制保证执行顺序。// GMM分组矩阵乘法的简化伪代码for(intg0;gn_groups;g){intexroute[g];// 路由到的专家编号float*ainputg*in_dim;// 当前组的输入激活float*wexpert_w[ex];// 专家权重矩阵float*bexpert_b[ex];// 专家偏置向量for(intr0;rout_dim;r){floataccb[r];for(intc0;cin_dim;c)acca[c]*w[r*in_dimc];output[g*out_dimr]act(acc);// 激活函数}}融合算子内部实现机制算子融合的判断条件基于两个核心指标中间张量的访存代价和kernel启动的调度代价。当两个相邻算子之间的中间结果可以完全容纳在片上存储如昇腾NPU的Unified Buffer中且该中间结果的生命周期不超过其父kernel的执行窗口时融合这两个算子能够消除至少一次全局存储的读写操作此时融合判定为正收益。收益评估方法是将融合前后的全局存储访问量做差值运算融合前的访存量等于各算子独立访存量之和加上中间张量的写入和读取量融合后的访存量仅包含输入和输出的端到端读写量。两者的差值即为融合节省的访存带宽当该差值超过阈值时判定融合可行。FusedAttention等融合kernel采用tile级拆分策略将注意力计算映射到昇腾NPU的向量计算单元上。拆分的基本单位是tile——一个固定大小的二维数据块尺寸由片上存储容量和计算单元的寄存器文件大小共同决定。在QKV投影阶段kernel将输入张量按tile维度拆分每个tile被加载到片上后完成Q、K、V三个投影的计算在缩放点积阶段Q的行tile与K的列tile在片上执行向量内积结果暂存于片上缓冲区Softmax归一化在片上完成分母的归约求和与逐元素除法加权求和阶段将归一化后的注意力权重与V的列tile执行向量外积累加。整个流程中中间结果始终驻留在片上存储中仅在最终输出时写回全局存储。ops-transformer仓库通过gitmodules引入catlass模板库建立了算子业务逻辑与底层计算模板的分层结构。catlass提供了面向昇腾NPU的矩阵乘法、向量归约、元素变换等基础计算模板这些模板经过手工调优充分利用了Cube单元和Vector单元的并行能力。融合算子在实现时将注意力计算中的矩阵乘法部分委托给catlass模板通过模板参数指定输入输出张量的布局、分块策略和累加精度而融合算子自身专注于控制流程编排、mask注入、Softmax归一化等业务逻辑。这种分层解耦使得catlass模板可以独立演进——当新的硬件特性或优化策略被纳入catlass时所有依赖catlass的融合算子无需修改即可自动受益。// tile级拆分的简化示意for(intti0;tin_tiles;ti){// 按行tile遍历load_tile(q,qiti*t_sz);// 加载Q tile到片上for(inttj0;tjn_tiles;tj){// 按列tile遍历load_tile(k,kjtj*t_sz);// 加载K tile到片上matmul_tile(q_tile,k_tile,s_tile);// catlass模板执行tile乘法scale_tile(s_tile,inv_d);// 缩放if(has_mask)apply_mask(s_tile);// 注入masksoftmax_tile(s_tile);// 片上Softmax归一化load_tile(v,vjtj*t_sz);// 加载V tile到片上matmul_tile(s_tile,v_tile,o_tile);// 加权求和store_tile(o,oitj*t_sz);// 写回输出tile}}MoE门控路由与top-k选择的数值实现moe子目录下的融合算子以GMM内核为核心围绕门控网络和top-k选择策略构建了一整套稀疏计算流程。门控网络的核心职责是将输入token映射到专家池中的特定子集这一映射过程在数学上表达为对专家集合的加权评分。对于每个输入token路由器通过一个线性层输出一个长度为专家总数的得分向量该向量经过Softmax归一化后得到每个专家被选中的概率分布。在top-k路由策略下系统从概率最高的k个专家中采样或直接选择作为当前token的执行目标其中k值通常取1或2。top-1路由对应专家选择完全由最高得分决定计算路径确定性最强top-2路由允许一个token触发两个专家的并行计算在专家覆盖率和计算量之间取得平衡。top-k选择在昇腾NPU上的实现需要兼顾排序效率和内存访问局部性。naive实现对长度为E的专家得分向量执行完整排序其复杂度为O(E log E)在专家数量较多时排序开销不可忽视。ops-transformer中的top-k实现采用分块比较策略先在片上缓存中找出当前tile内的局部top-k值再通过并行归约将各tile的局部结果合并为全局top-k专家集合。这种两阶段策略将排序空间压缩到以tile为单位的局部数据上大幅降低了比较操作的次数。实现中还需要处理的一个重要边界条件是当多个专家得分相等时的仲裁策略通常以专家编号作为次级排序键来保证结果确定性。Softmax数值稳定性与log-sum-exp技巧SMLAN与SMLA两个融合系列在Softmax数值实现上存在关键差异这一差异直接影响算子在异常输入下的鲁棒性表现。标准Softmax在数学上对输入向量中各元素执行指数函数映射后归一化当输入元素的绝对值较大时指数运算的输出值可能超出浮点表示范围而导致上溢。在FP16精度下指数函数在输入值超过约11时开始溢出在BF16精度下上溢阈值提升至约89但下溢问题在接近零的负数区域依然存在。SMLAN系列在处理无mask场景时假设输入分布可控但在实际部署中经过多层Transformer堆叠后的注意力分数可能累积放大数值稳定性风险随之增加。log-sum-exp技巧是解决Softmax数值稳定性问题的标准方法其核心思想是在指数运算前将输入向量平移一个与最大值相关的偏移量。在数学上对输入向量x的每个元素减去max(x)等价于将指数求和项乘以exp(-max(x))由于exp(-max(x))不大于1该乘法不会导致上溢。标准化结果则通过在分母的log-sum-exp项上加上max(x)来还原。ops-transformer在SMLA系列中引入了这一技巧通过在kernel启动时执行一次并行归约求出向量最大值并将该值广播到所有计算线程中作为平移基准。在昇腾NPU的Vector单元上实现log-sum-exp需要特别关注归约过程的同步开销——如果归约结果需要通过全局存储中转才能分发给所有线程则一次额外的全局读写会抵消log-sum-exp带来的数值收益。ops-transformer通过将归约过程完全保留在片上执行来解决这一问题利用Vector单元的 warp级归约指令求出max值后通过寄存器广播将结果直接传递给同warp内的其他线程整个过程无需离开片上存储完成。catlass模板与ops-transformer的协同架构catlass作为昇腾NPU的底层计算模板库承担着将硬件计算能力抽象为可复用接口的职责。catlass的核心抽象是tile级GEMM模板——它将矩阵乘法拆解为一系列tile级别的分块操作每个tile在片上完成计算后通过异步DMA流写回全局存储同时启动下一批tile的加载以隐藏数据传输延迟。catlass定义了tile循环的通用骨架外层循环遍历输出空间内层循环遍历分块索引调度逻辑负责将数据预取与计算重叠。对于调用方而言只需提供tile级别的计算回调函数而无需关心tile调度和存储管理的细节。ops-transformer在catlass的tile骨架基础上进行定制化修改以满足融合算子的业务需求。定制化修改主要体现在三个方面。第一输入张量布局的适配catlass原生支持行优先和列优先两种数据布局而attention计算中Q、K、V通常以head维度优先的布局组织ops-transformer通过在catlass模板的tile加载回调中插入布局转置逻辑使得融合算子可以使用原生注意力布局而无需预先对输入张量进行重排。第二分块策略的微调catlass的标准GEMM分块基于输出行块和权重列块进行划分而attention中的加权求和阶段需要按注意力权重的列维度与V的列块对齐ops-transformer在catlass的tile索引映射表中增加了列对齐约束使得不同阶段的tile划分在边界处能够精确对齐。第三中间结果的暂存策略catlass的GEMM模板假设每个tile的计算结果直接写回或参与下一层计算而融合算子需要在Softmax阶段保留完整的注意力分数矩阵作为中间结果ops-transformer通过扩展catlass的片上缓冲区管理接口为中间矩阵分配专用空间并在其生命周期结束时显式释放从而在利用catlass计算效率的同时保留了中间数据的可用性。MC2多通道复制的内存复用机制mc2子目录下的Multi-Channel Copy算子在多token并行生成场景中承担着关键的内存带宽优化角色。自回归解码的投机采样或推测解码范式需要模型在单步推理中生成并验证多个候选token这些候选token在注意力计算时共享相同的KV缓存内容直到某个token的验证失败为止。在此之前所有候选分支的KV状态完全一致因此各分支重复读取同一份KV缓存会产生大量冗余的全局存储访问。MC2算子针对这一冗余模式设计了一次读取、多次广播的执行语义从根本上消除了跨通道的重复读取开销。MC2的内存复用机制建立在通道维度拆分和片上广播两个技术支柱之上。在通道维度拆分层面MC2将多通道复制操作拆解为源数据的单次加载和目标通道的多次写入两个子过程。源数据通常是KV缓存的某一段在kernel启动后由所有工作线程协同读取一次将其放入片上的广播缓冲区随后各通道的工作组分别从该广播缓冲区中读取数据并写入各自的目标地址。由于源数据只被加载一次多个目标通道的总读取带宽等于单通道读取带宽总写入带宽等于通道数乘以单通道写入带宽在通道数较多时这种不对称的带宽消耗模式能够有效降低推理延迟。MC2在任务流编排中的位置决定了其对整体推理吞吐量的影响程度。在MoE架构的多token推理流水线上MC2通常作为KV缓存分发的前置阶段出现为后续GMM在各候选通道中的分组计算准备输入数据。MC2的kernel执行时间与KV缓存的大小和候选通道数成正比但与序列长度无关——这意味着当候选token数量增加时MC2的额外开销以常数级增长而每个token独立读取KV缓存的方式其额外开销则以线性级增长。在候选token数达到8至16的推测解码场景下这一差异尤为显著。此外MC2算子的融合设计与昇腾NPU的统一缓存架构形成协同——源数据从Global Memory加载到片上后广播写入各通道的过程中数据始终保留在统一缓存层级中减少了缓存层级间的回写和重新加载次数。效率对比维度使用前使用后差异来源QKV投影至Softmax的中间数据搬移Q、K、V三个张量各写入全局存储一次三个张量驻留片上无全局存储写入融合消除了3次中间张量的全局写入与读取Softmax归一化的kernel启动次数Softmax作为独立kernel启动每次启动含调度开销Softmax在融合kernel内完成无额外启动开销融合减少了至少1次核函数启动的调度延迟MoE分组乘法的权重加载效率全量加载所有专家权重后按路由掩码选择按路由分组仅加载被选中专家的权重分块路由稀疏性决定专家利用率越低节省越多多Token验证的KV缓存读取次数每个候选token独立读取完整KV缓存共享KV缓存在片上广播至各候选通道MC2广播替代重复读取节省量与候选token数成正比结尾ops-transformer算子库围绕Transformer模型的核心计算模块在CANN框架上构建了attention融合、GMM分组乘法、MC2多通道复制等专用算子实现。每类算子针对昇腾NPU的存储层次和计算单元特性进行了适配通过算子融合消除中间访存、通过分组策略匹配稀疏路由、通过广播机制压缩冗余读取。仓库的分层架构将业务逻辑与底层计算模板解耦catlass模板库的引入使得高性能矩阵乘法等基础计算的优化成果能够被各类融合算子复用。仓库地址https://atomgit.com/cann/ops-transformer