CANN ops-transformer：Transformer 算子全家桶一览-开发者社区

个人主页：ujainu

文章目录

前言
- 先搞清楚它都装了些什么
- - Attention 家族
  - MoE 类
  - 位置编码与归一化
  - MC2 通算融合类
- 这些算子靠谁干活
- 谁在用它
- 图里看一眼位置

前言

说人话：ops-transformer 就是昇腾 CANN 生态里专门为大模型准备的那套算子工具箱。你跑 GPT、DeepSeek、Qwen 这类 Transformer 架构的模型，在昇腾 NPU 上实际执行的就是它里面的那些算子——FlashAttention、MoE 路由、RMSNorm、旋转位置编码……全在这一堆里。

为什么需要单独搞一个仓库？因为 Transformer 模型的计算模式和传统的 CNN 完全不一样。CNN 反复做卷积池化，算子种类少但调用频次高；Transformer 里 attention 要查表、MoE 要选专家、KV Cache 要动态管理，每个环节都是一套独立的计算逻辑。ops-math 和 ops-nn 那两个基础算子库能搞定通用场景，但 Transformer 的那些"偏门操作"得有人专门写。ops-transformer 就是干这个的。

再往大了说，大模型推理的瓶颈几乎全卡在 Transformer 特有的那些操作上：长序列的 Attention 是显存杀手，MoE 的 expert 路由是多卡通信的重灾区，KV Cache 的动态扩缩是推理框架的头疼事。这些问题的解法不是通用矩阵乘能覆盖的，需要针对每个场景单独做算子级的优化。ops-transformer 在整个 CANN 生态里的定位，就是那个"专治 Transformer 各种疑难杂症"的专科大夫——别的算子库看不了的病，找它就对了。

先搞清楚它都装了些什么

这个仓库目前覆盖了 15 个以上的算子，粗略分四大类来看：

Attention 家族

这是 Transformer 最核心的计算，也是显存占用的大头。标准 Attention 的复杂度是 O(n²)，序列一长就炸。ops-transformer 里一口气给了好几个变体，每个解决不同场景的问题。

# 最经典的 FlashAttentionScore# 把 Q/K 切成小块送进 SRAM 算，算完就丢，不用把整个 S 矩阵塞进片外内存# 和原版 FlashAttention 论文思路一致，但针对昇腾 NPU 的 Cube 单元做了适配out=flash_attention_score(q,k,v,mask)

FlashAttentionScore 的核心思路是分块计算（tiling）。传统 Attention 需要先算出完整的 QK^T 注意力矩阵，再对 V 做加权求和，中间那个 S 矩阵的大小是 seq_len × seq_len，seq_len 到 128K 的时候这个矩阵根本放不下。FlashAttention 把 Q 和 K 都切成小块，每块在片上 SRAM 里算完就写回结果，中间的 S 矩阵不用存到 HBM，显存占用从 O(n²) 降到了 O(n)。

# SparseFlashAttention——长序列场景用的# 有些 token 的 attention 分数极低，直接跳过不算，省掉大量无效计算# 背后是稀疏注意力模式：local + global + random 三种窗口混合out=sparse_flash_attention(q,k,v,sparsity_mask)

SparseFlashAttention 在 FlashAttention 的基础上再加一层稀疏过滤。大模型推理到 128K 甚至 256K 序列长度时，即便有了 FlashAttention 的分块优化，计算量仍然是 O(n²)，算不完。稀疏注意力的观察是：大多数 token 对之间的注意力分数接近零，真正有意义的连接只是少数。SparseFlashAttention 根据一个稀疏掩码跳过那些"不值得算"的 token 对，把计算量从 O(n²) 压到 O(n × k)，k 是每个 token 实际关注的目标数量。

# GatherPAKVCache——推理时把新生成的 KV 拼进缓存# 不用每次都重新算，只追加增量部分# PAK = Page Attention Key，配合 PagedAttention 的分页管理new_cache=gather_pakv_cache(kv_cache,new_k,new_v,page_table)

GatherPAKVCache 解决的是自回归推理的缓存管理问题。每生成一个新 token，就会产生一组新的 K 和 V 向量，需要拼到已有的 KV Cache 后面。听起来简单，但 KV Cache 在显存里不一定是连续存放的——PagedAttention 把它分成固定大小的 page，离散分布在显存里。GatherPAKVCache 的工作就是根据 page_table 把新的 KV 数据写到正确的 page 位置上，同时维护索引，让后续的 Attention 读取时能正确拼出完整的 K 和 V 序列。

MoE 类

混合专家模型（如 DeepSeek-V3）每个 token 要选几条专家去计算，这个路由和调度逻辑需要专门的算子来支撑。

# MoE 路由：每个 token 选 top-k 个专家，把 token 分发过去# gate_logits 是路由网络的输出，shape 为 [num_tokens, num_experts]tokens_per_expert=moe_compute_expert_tokens(hidden,gate_logits,top_k=8)

MoEComputeExpertTokens 的原理可以拆成三步：第一步，路由网络（一个小的线性层 + softmax）为每个 token 产出对所有专家的偏好分数；第二步，取 top-k，选出分数最高的 k 个专家；第三步，把 token 按照"要去哪个专家"重新排列，同一专家的 token 连续排放，方便后续的批量矩阵乘。第三步的排列操作是性能关键——如果 token 顺序混乱，每个专家拿到的输入就是不连续的内存片段，Cube 单元的利用率会骤降。

# MoE AlltoAll 通信：把 token 按专家归属分发到对应卡上# 多卡场景下，token 可能要去其他 NPU 上的专家# 这个操作和通信绑在一起，是 MC2 融合算子的典型场景recv_tokens=moe_alltoall_send_recv(send_tokens,expert_map,comm_handle)

位置编码与归一化

KVRMSNormRoPECache 把旋转位置编码（RoPE）和 RMSNorm 融合在一起算，一趟完成两件事。

# RoPE + RMSNorm 融合算子# 单独写就是两次搬运，合并算子省掉中间数据的来回拷贝# RoPE: 对 Q/K 向量按位置旋转，注入位置信息# RMSNorm: 对隐层向量做归一化，稳定训练out=kv_rmsnorm_rope_cache(x,rope_cos,rope_sin,gamma)

为什么要把这两个操作合成一个算子？因为在大模型的推理循环里，RMSNorm 和 RoPE 是紧挨着执行的——隐层输出先过 RMSNorm 归一化，再过 RoPE 注入位置编码，然后才送进 Attention。如果分开算，RMSNorm 的输出要先写回 HBM，RoPE 再从 HBM 读进来，一来一回就是两次显存搬运。融合成一个算子后，RMSNorm 的结果留在片上 SRAM 里直接喂给 RoPE，省掉一次读一次写，对推理延迟的影响不小。

# 单独的 SwiGLU 激活——FFN 层的核心# SwiGLU = Swish(xW1) ⊗ xW2，比 ReLU 效果好，但多一个矩阵乘# 在 LLaMA/Qwen/DeepSeek 的 FFN 里都用它out=swiglu(x,w1,w2)# out = swish(x @ w1) * (x @ w2)

SwiGLU 虽然归在 norm/激活类里，但它本质是 FFN 层的前半段。标准 FFN 是两层线性变换夹一个 ReLU，SwiGLU 改成了两个门控分支：一个分支做 Swish 激活，另一个分支不做激活，两个分支逐元素相乘。多了一个分支意味着多一次矩阵乘，但换来的是模型表达力的提升。ops-transformer 把 SwiGLU 封装成独立算子，而不是让框架自己拼两个 matmul + 一个 mul，是因为在昇腾 NPU 上可以把两次矩阵乘和数据搬运合并调度，减少指令下发开销。

MC2 通算融合类

这是 ops-transformer 里比较特别的一块。MatmulAlltoAll、AttentionToFFN、FFNToAttention 这几个算子把矩阵乘法和通信操作绑在了一起，算着算着就把数据发走了，通信和计算重叠执行。

MC2 的思路是：大模型多卡推理时，通信等待是很大的开销。Attention 算完要把结果发给下一组卡跑 FFN，FFN 算完又要发回来。如果算完再通信，通信期间 NPU 就在干等。MC2 融合算子的做法是——矩阵乘算出一块结果就立刻启动通信发送，同时继续算下一块，计算和通信像流水线一样交替推进。

# 看看仓库目录结构，大致长这样ops-transformer/ ├── attention/# FlashAttention 等变体├── moe/# MoE 路由与分发├── norm/# RMSNorm 融合算子├── rope/# 旋转位置编码├── mc2/# 通算融合算子└── swiglu/# SwiGLU 激活相关

# 用 msprof 看一下算子在 NPU 上的执行情况# 可以清楚看到 FlashAttention 的 Cube 利用率和 MC2 的通信计算重叠效果msprof--application="python infer.py"--output="./profiling_data"# 生成报告后打开 msprof 可视化工具查看# 关注指标：Cube 单元利用率、HBM 带宽、通信占比

这些算子靠谁干活

ops-transformer 自己不生产原材料，它盖房子得有砖。往上追溯，它依赖两个关键的上游仓库。

opbase是所有算子仓库的公共底座，提供头文件、数据结构定义、调度框架这些基础的东西。你可以理解为 opbase 定义了算子的"施工标准"——输入输出长什么样、怎么在 NPU 上排队执行，这些规矩都在 opbase 里定好。ops-transformer 里每个算子的底层都依赖 opbase 提供的基础组件。

// opbase 提供的算子基类，ops-transformer 的每个算子都继承它// 下面是简化的 C++ 伪代码，展示算子注册和执行框架classOpBase{public:virtualStatusInit(constOpParam&param)=0;// 初始化：分配资源、检查形状virtualStatusExecute(constTensorList&in,TensorList&out)=0;// 执行：调 Ascend C kernelvirtualStatusDestroy()=0;// 清理};// ops-transformer 里的 FlashAttention 算子继承 OpBaseclassFlashAttentionScore:publicOpBase{StatusInit(constOpParam&param)override{seq_len_=param.Get("seq_len");head_dim_=param.Get("head_dim");returnStatus::OK();}StatusExecute(constTensorList&in,TensorList&out)override{// 调用 Ascend C 编写的 tiling kernelreturnLaunchFlashAttentionKernel(in,out,tiling_data_);}};

ascend-boost-comm是算子公共平台中间件，ops-transformer 里的 MC2 通算融合算子要用到它。通信和计算要绑在一起跑，需要一个中间层来调度两边的工作节奏，这就是 ascend-boost-comm 干的事——它不负责具体通信（那是 hccl 的事），而是做一个 M×N 的算子复用调度。具体说，当有 M 种计算模式和 N 种通信模式排列组合时，ascend-boost-comm 提供统一的调度接口，避免为每种组合都写一套融合逻辑。

# MC2 融合算子的调用方式——通信句柄由 ascend-boost-comm 提供importacl# Ascend Computing Language# 初始化通信资源comm_handle=ascend_boost_comm.init(rank_id=0,world_size=8)# 调用 AttentionToFFN：Attention 计算结果边算边发给 FFN 所在的卡# 内部流程：matmul 出一块 → 立刻 alltoall 发走 → 继续 matmul 下一块ffn_input=attention_to_ffn(attn_output,comm_handle)# FFN 算完，再发回 Attention 卡next_attn_input=ffn_to_attention(ffn_output,comm_handle)

谁在用它

ops-transformer 在 CANN 五层架构里属于第 2 层（AOL 算子库层）。从上到下看这五层：最上面是应用层（你的模型脚本），第 4 层是框架适配层（PyTorch/MindSpore 的昇腾适配），第 3 层是加速库层（ATB 在这里），第 2 层就是算子库层（ops-transformer、ops-math、ops-nn 都在这），第 1 层是昇腾 CANN 的底层运行时和驱动。ops-transformer 处在"承上启下"的关键位置——往上给 ATB 和加速库提供高性能算子实现，往下调用 opbase 的框架和昇腾 NPU 的硬件能力。

最大的下游是ATB（ascend-transformer-boost，Transformer 加速库）。ATB 提供了三层架构——基础原生算子、图算子机制、插件机制——其中"基础原生算子"这一层本质上就是对 ops-transformer 算子的高层封装。你在 PyTorch 里跑大模型，ATB 在中间把整个 Transformer 前向传播编排好，底层实际执行的就是 ops-transformer 的 FlashAttention、SwiGLU、RMSNorm 这些算子。

# ATB 的图算子机制：把多个 ops-transformer 算子编排成一张子图# 比如一个完整的 Decoder Layer 的前向传播classDecoderLayerGraph:defbuild(self,graph_builder):# Step 1: RMSNorm + RoPE 融合norm_rope=graph_builder.add_op("kv_rmsnorm_rope_cache",inputs=[x,cos,sin,gamma])# Step 2: FlashAttentionattn=graph_builder.add_op("flash_attention_score",inputs=[norm_rope,kv_cache])# Step 3: MC2 融合——Attention 到 FFNffn_in=graph_builder.add_op("attention_to_ffn",inputs=[attn,comm_handle])# Step 4: SwiGLU + 矩阵乘ffn_out=graph_builder.add_op("swiglu",inputs=[ffn_in,w1,w2])# Step 5: MC2 融合——FFN 回到 Attentionreturngraph_builder.add_op("ffn_to_attention",inputs=[ffn_out,comm_handle])

再往上就是cann-recipes-infer和cann-recipes-train这两个配方仓库。它们是面向开发者的端到端方案，把 ATB + ops-transformer + 框架适配全部串起来。你想在昇腾上部署 DeepSeek 推理，直接拿 cann-recipes-infer 的脚本跑就行，ops-transformer 的算子在底层默默干活。

# cann-recipes-infer 的典型使用方式# 一条命令拉起 DeepSeek 推理，底层自动调 ops-transformer 的算子bashrun_inference.sh\--modeldeepseek-v3\--npu0,1,2,3,4,5,6,7\--seq-length 128k\--batch-size1# 配置文件里可以指定使用哪些算子变体# 比如 Attention 可以选 flash / sparse / paged 三种模式

还有一个容易被忽略的关联仓库：graph-autofusion。它是算子自动融合框架，能自动把 ops-transformer 里的多个算子合并成一个 SuperKernel 执行。比如把 RMSNorm + RoPE + KVCache 拼接三个算子合成一步，少两次中间数据的搬运。graph-autofusion 和 ops-transformer 里 KVRMSNormRoPECache 那种手动融合的区别在于：手动融合需要开发者提前知道哪些算子要合在一起写代码实现；自动融合是框架在运行时分析计算图，自动发现可以合并的算子组合。两者互补——手动融合覆盖已知的固定模式，自动融合捕捉动态的优化机会。

# graph-autofusion 的融合规则配置（示意）# 告诉框架哪些算子模式可以自动合并成 SuperKernelfusion_rules:-name:"norm_rope_fusion"pattern:["rmsnorm","rope"]# 顺序执行的算子模式action:"merge_to_super_kernel"# 合并为一个 SuperKernel-name:"attn_residual_fusion"pattern:["flash_attention","residual_add"]action:"merge_to_super_kernel"

图里看一眼位置

┌──────────────────────────┐ │ cann-recipes-infer/train │ │ (端到端配方) │ └────────────┬─────────────┘ │ ┌────────────▼─────────────┐ │ ATB (Transformer 加速库) │ └────────────┬─────────────┘ │ ┌─────────────────▼──────────────────┐ │ ops-transformer (本文主角) │ │ 第2层 AOL 算子库 │ └───┬──────────────┬────────────┬────┘ │ │ │ ┌───────▼──────┐ ┌────▼─────┐ ┌───▼──────────┐ │ opbase │ │ascend- │ │graph-auto- │ │ (算子底座) │ │boost- │ │fusion(融合) │ │ │ │comm │ │ │ └──────────────┘ └──────────┘ └──────────────┘

如果你在做大模型相关的昇腾开发，ops-transformer 基本绕不开——不管你走 ATB 的上层封装还是直接调算子，它都在那里。仓库地址贴一下：https://atomgit.com/cann/ops-transformer ，clone 下来看看源码目录结构，对照上面提到的算子分类走一遍，比看任何文档都快。