个人主页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 单元做了适配outflash_attention_score(q,k,v,mask)FlashAttentionScore 的核心思路是分块计算tiling。传统 Attention 需要先算出完整的 QK^T 注意力矩阵再对 V 做加权求和中间那个 S 矩阵的大小是 seq_len × seq_lenseq_len 到 128K 的时候这个矩阵根本放不下。FlashAttention 把 Q 和 K 都切成小块每块在片上 SRAM 里算完就写回结果中间的 S 矩阵不用存到 HBM显存占用从 O(n²) 降到了 O(n)。# SparseFlashAttention——长序列场景用的# 有些 token 的 attention 分数极低直接跳过不算省掉大量无效计算# 背后是稀疏注意力模式local global random 三种窗口混合outsparse_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_cachegather_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_expertmoe_compute_expert_tokens(hidden,gate_logits,top_k8)MoEComputeExpertTokens 的原理可以拆成三步第一步路由网络一个小的线性层 softmax为每个 token 产出对所有专家的偏好分数第二步取 top-k选出分数最高的 k 个专家第三步把 token 按照要去哪个专家重新排列同一专家的 token 连续排放方便后续的批量矩阵乘。第三步的排列操作是性能关键——如果 token 顺序混乱每个专家拿到的输入就是不连续的内存片段Cube 单元的利用率会骤降。# MoE AlltoAll 通信把 token 按专家归属分发到对应卡上# 多卡场景下token 可能要去其他 NPU 上的专家# 这个操作和通信绑在一起是 MC2 融合算子的典型场景recv_tokensmoe_alltoall_send_recv(send_tokens,expert_map,comm_handle)位置编码与归一化KVRMSNormRoPECache 把旋转位置编码RoPE和 RMSNorm 融合在一起算一趟完成两件事。# RoPE RMSNorm 融合算子# 单独写就是两次搬运合并算子省掉中间数据的来回拷贝# RoPE: 对 Q/K 向量按位置旋转注入位置信息# RMSNorm: 对隐层向量做归一化稳定训练outkv_rmsnorm_rope_cache(x,rope_cos,rope_sin,gamma)为什么要把这两个操作合成一个算子因为在大模型的推理循环里RMSNorm 和 RoPE 是紧挨着执行的——隐层输出先过 RMSNorm 归一化再过 RoPE 注入位置编码然后才送进 Attention。如果分开算RMSNorm 的输出要先写回 HBMRoPE 再从 HBM 读进来一来一回就是两次显存搬运。融合成一个算子后RMSNorm 的结果留在片上 SRAM 里直接喂给 RoPE省掉一次读一次写对推理延迟的影响不小。# 单独的 SwiGLU 激活——FFN 层的核心# SwiGLU Swish(xW1) ⊗ xW2比 ReLU 效果好但多一个矩阵乘# 在 LLaMA/Qwen/DeepSeek 的 FFN 里都用它outswiglu(x,w1,w2)# out swish(x w1) * (x w2)SwiGLU 虽然归在 norm/激活类里但它本质是 FFN 层的前半段。标准 FFN 是两层线性变换夹一个 ReLUSwiGLU 改成了两个门控分支一个分支做 Swish 激活另一个分支不做激活两个分支逐元素相乘。多了一个分支意味着多一次矩阵乘但换来的是模型表达力的提升。ops-transformer 把 SwiGLU 封装成独立算子而不是让框架自己拼两个 matmul 一个 mul是因为在昇腾 NPU 上可以把两次矩阵乘和数据搬运合并调度减少指令下发开销。MC2 通算融合类这是 ops-transformer 里比较特别的一块。MatmulAlltoAll、AttentionToFFN、FFNToAttention 这几个算子把矩阵乘法和通信操作绑在了一起算着算着就把数据发走了通信和计算重叠执行。MC2 的思路是大模型多卡推理时通信等待是很大的开销。Attention 算完要把结果发给下一组卡跑 FFNFFN 算完又要发回来。如果算完再通信通信期间 NPU 就在干等。MC2 融合算子的做法是——矩阵乘算出一块结果就立刻启动通信发送同时继续算下一块计算和通信像流水线一样交替推进。# 看看仓库目录结构大致长这样ops-transformer/ ├── attention/# FlashAttention 等变体├── moe/# MoE 路由与分发├── norm/# RMSNorm 融合算子├── rope/# 旋转位置编码├── mc2/# 通算融合算子└── swiglu/# SwiGLU 激活相关# 用 msprof 看一下算子在 NPU 上的执行情况# 可以清楚看到 FlashAttention 的 Cube 利用率和 MC2 的通信计算重叠效果msprof--applicationpython infer.py--output./profiling_data# 生成报告后打开 msprof 可视化工具查看# 关注指标Cube 单元利用率、HBM 带宽、通信占比这些算子靠谁干活ops-transformer 自己不生产原材料它盖房子得有砖。往上追溯它依赖两个关键的上游仓库。opbase是所有算子仓库的公共底座提供头文件、数据结构定义、调度框架这些基础的东西。你可以理解为 opbase 定义了算子的施工标准——输入输出长什么样、怎么在 NPU 上排队执行这些规矩都在 opbase 里定好。ops-transformer 里每个算子的底层都依赖 opbase 提供的基础组件。// opbase 提供的算子基类ops-transformer 的每个算子都继承它// 下面是简化的 C 伪代码展示算子注册和执行框架classOpBase{public:virtualStatusInit(constOpParamparam)0;// 初始化分配资源、检查形状virtualStatusExecute(constTensorListin,TensorListout)0;// 执行调 Ascend C kernelvirtualStatusDestroy()0;// 清理};// ops-transformer 里的 FlashAttention 算子继承 OpBaseclassFlashAttentionScore:publicOpBase{StatusInit(constOpParamparam)override{seq_len_param.Get(seq_len);head_dim_param.Get(head_dim);returnStatus::OK();}StatusExecute(constTensorListin,TensorListout)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_handleascend_boost_comm.init(rank_id0,world_size8)# 调用 AttentionToFFNAttention 计算结果边算边发给 FFN 所在的卡# 内部流程matmul 出一块 → 立刻 alltoall 发走 → 继续 matmul 下一块ffn_inputattention_to_ffn(attn_output,comm_handle)# FFN 算完再发回 Attention 卡next_attn_inputffn_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 的硬件能力。最大的下游是ATBascend-transformer-boostTransformer 加速库。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_ropegraph_builder.add_op(kv_rmsnorm_rope_cache,inputs[x,cos,sin,gamma])# Step 2: FlashAttentionattngraph_builder.add_op(flash_attention_score,inputs[norm_rope,kv_cache])# Step 3: MC2 融合——Attention 到 FFNffn_ingraph_builder.add_op(attention_to_ffn,inputs[attn,comm_handle])# Step 4: SwiGLU 矩阵乘ffn_outgraph_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_fusionpattern:[rmsnorm,rope]# 顺序执行的算子模式action:merge_to_super_kernel# 合并为一个 SuperKernel-name:attn_residual_fusionpattern:[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 下来看看源码目录结构对照上面提到的算子分类走一遍比看任何文档都快。