ops-transformer算子库深度解读：FlashAttention与MoE在昇腾CANN上的实现思路
前言
跑一个千亿参数的大模型推理，最让人头疼的不是模型本身有多大，而是 attention 那一层的计算把显存和带宽吃得干干净净。你在昇腾NPU上搭一套推理服务，如果直接用框架自带的算子跑注意力机制，很快就会发现整个流程被 attention 侧的内存搬运卡住了。CANN 社区在 ops-transformer 仓库里给出了一整套针对 Transformer 结构的进阶算子实现，把 FlashAttention、MoE 路由计算、MC2 通算融合这些大模型里最吃资源的部分做了硬件级优化。这篇文章会把这套算子库的几个核心模块拆开来看，弄清楚它们到底在做什么，以及为什么这样设计。
注意力机制的显存墙
Transformer 模型的核心计算路径是 QKV 注意力。给定输入序列长度 n 和模型隐藏维度 d，标准注意力的做法是先把 Q、K、V 三个矩阵算出来，Q 乘 K 的转置得到一个 n 乘 n 的注意力分数矩阵，对这个矩阵做 softmax，再乘 V 得到输出。这套流程在 CPU 上跑没问题，在 GPU 上勉强能撑住中小模型，但到了大模型场景，n 乘 n 这个中间矩阵会直接撑爆显存。序列长度 8192 的时候，单层 attention 的分数矩阵就需要好几个 GB 的显存来存放，而且这个矩阵要读写各一次，带宽压力巨大。
昇腾达芬奇架构的 Cube 单元擅长做矩阵乘法，但它的片上存储（L0A/L0B/L0C）容量有限。如果attention计算过程中频繁地把中间结果搬进搬出 Global Memory，Cube 单元再快也是白搭，因为带宽成了瓶颈。ops-transformer 仓库里的 flash_attn 目录专门解决这个问题。它的核心思路是：不再一次性算完整个 n 乘 n 的注意力矩阵，而是把 Q 和 K 按行和列切成小块，在片上存储里逐块计算 softmax，算完一块就把这块的结果累加到输出里，小块的中间结果始终留在片上，不需要写回 Global Memory。
这样做的好处是显存占用从 O(n^2) 降到了 O(n)，因为不需要存完整的注意力分数矩阵了。带宽消耗也大幅下降，因为中间结果不再反复搬运。代价是计算逻辑变复杂了——你得自己管理分块策略、在线 softmax 的数值稳定性，以及反向传播时分块梯度的重组。这些都是框架自带算子不需要操心的事情。
FlashAttention 的分块计算
理解 flash_attn 的实现，要先搞清楚两个概念：tiling 和 online softmax。
Tiling 就是把大矩阵切成小块（tile）。昇腾 NPU 的 Cube 单元一次能处理的矩阵大小受限于片上存储容量，所以必须把 Q、K、V 按合适的粒度切块，一块一块地喂给 Cube 单元。切块的策略不是随便切——块太大会溢出片上存储，块太小则 Cube 单元的并行度不够，计算效率低。
Online softmax 是另一个关键技巧。标准 softmax 需要两遍扫描：第一遍找最大值，第二遍算指数和归一化。但分块计算时你不可能先把所有块的分数收集起来再统一做 softmax，因为那又回到了存完整分数矩阵的老路。online softmax 的做法是在处理每个块的时候，维护一个"当前全局最大值"和一个"当前指数和"，每处理一个新块就更新这两个值，同时回溯修正之前已算出的块的输出。这样一遍扫描就能得到与全局 softmax 等价的结果。
ops-transformer 的 flash_attn 算子用 Ascend C 语言实现，直接操作 Cube 向量计算单元和片上存储。一段简化后的核心逻辑大致是这样的：
// flash_attn 分块计算的核心循环// x 是 Q 的一行块，y 是 K 的一列块，s 是注意力分数for(inti=0;i<q_blocks;i++){floatm=-INF;// 当前最大值floatl=0.0f;// 当前指数和// 先把输出块 o 清零for(intj=0;j<k_blocks;j++){// s = x[i] * y[j]^T，在 Cube 上算，结果留在片上s=matmul(x[i],y[j]);// online softmax: 逐块更新最大值和指数和floatm_new=max(m,max(s));floatexp_old=exp(m-m_new);floatexp_new=exp(max(s)-m_new);l=l*exp_old+sum(exp_new*s);m=m_new;// 把修正后的结果写回输出o[i]=o[i]*exp_old+exp(s-m)*v[j];}}这样写是因为 NPU 的 Cube 单元只接受固定大小的矩阵乘法调用，分块大小必须与硬件参数对齐。online softmax 的更新逻辑必须内嵌在循环里，不能拆成单独的 kernel，否则中间结果就要写回 Global Memory。整个循环的控制流在 Ascend C 里用 local memory 做 buffer，保证数据始终在片上流动。
MoE 路由计算的并行困境
混合专家模型（MoE）是近年来大模型 scaling 的一条主流路线。它的思路是把 FFN 层拆成多个"专家"，每个 token 根据路由器的打分只激活其中少数几个专家。这种架构能以接近常量的计算开销增加模型参数量，让模型在不增加推理计算量的前提下获得更强的表达能力。
但 MoE 的算子实现有一个天然难题：不同 token 被分配到不同专家，而每个专家的处理需要把属于它的 token 收集到一起做矩阵乘法。这个"收集"操作涉及大量不连续的内存访问和 all-to-all 通信，在分布式场景下更麻烦。ops-transformer 的 moe 目录里有一系列算子专门处理这个流程：moe_gating_top_k 做路由打分和 top-k 选择，moe_token_permute 把 token 按专家重新排列，moe_finalize_routing 把专家计算完的结果拼回去。
这套流程的瓶颈在通信而不是计算。permute 操作本质是对索引数组做 gather，每个 token 要从原始位置搬到一个按专家分组的新位置。在多卡场景下，不同专家可能分布在不同卡上，就变成了 all-to-all 通信。昇腾 NPU 的 HCCL 通信库支持 all-to-all，但通信延迟取决于数据量和对齐方式。如果 permute 之后的数据在内存中是零散的，还得先做 pack 再发，多了一次内存拷贝。
ops-transformer 对这部分的处理是：把 gating、permute、dispatch 这几个步骤拆成独立的算子，每个算子只负责一件事，但它们之间通过 shared memory 共享路由索引，避免重复计算。token 的重排用 SIMT（单指令多线程）模式处理，每条线程负责一批 token 的索引映射。
# moe 路由计算的调用流程（示意）# q 是输入张量，w_gate 是路由权重gate_logits=ops.matmul(q,w_gate)# top_k 选出每个 token 去哪几个专家topk_idx,topk_val=ops.topk(gate_logits,k=2)# 把 token 按专家重排permuted=ops.moe_token_permute(q,topk_idx)# 每个专家独立做 FFNexpert_out=ops.grouped_matmul(permuted,expert_weights,topk_idx)# 把结果拼回原始顺序out=ops.moe_finalize_routing(expert_out,topk_idx)把路由拆成独立算子的原因是每个步骤的数据分布特征不一样。gating_top_k 是密集的矩阵乘法加 top-k，适合 Cube 单元。token_permute 是不连续内存访问，适合 Vector 单元。finalize_routing 需要做 scatter 加归一化，又是另一种访问模式。硬塞进一个融合算子里反而会因为数据布局不匹配而降低效率。分而治之，让每个算子跑在最适合自己的硬件单元上。
MC2：通信与计算的融合
MC2（Matmul-Communication-Compute fusion）是 ops-transformer 仓库里另一个重要的算子类别。在大模型分布式训练和推理中，矩阵乘法做完之后经常紧跟一个通信操作（all-reduce、reduce-scatter、all-to-all），之后再进入下一层的计算。传统做法是这三个阶段串行执行：先算矩阵乘，等算完了发起通信，通信完成了再算下一层。
MC2 的思路是把矩阵乘和通信融合成一个算子。矩阵乘的输出不需要写回 Global Memory，直接作为通信操作的输入，通信完成后再进入下一层的计算。这样做省掉了矩阵乘输出到 Global Memory 的写入和通信输入从 Global Memory 的读取——两次内存搬运被消掉了。
ops-transformer 的 mc2 目录里有 matmul_all_reduce、matmul_reduce_scatter、allto_all_matmul、attention_to_ffn 等算子。以 matmul_all_reduce 为例，它把线性层的矩阵乘和跨卡 all-reduce 融合在一起：
// matmul_all_reduce 的计算流水线// a 是输入，b 是权重// 算完 matmul 的结果不落地，直接进入 all-reducevoidfused_mm_ar(void*a,void*b,void*c,hccl_ctx ctx){// Cube 做 matmul，结果留在片上matmul(a,b,local_buf);// 把片上结果直接送给 HCCL，不写 Global Memoryhccl_all_reduce(local_buf,c,count,dtype,ctx);// all-reduce 的结果直接作为后续计算的输入}不落地的设计目的是让计算和通信的时间重叠起来。NPU 有独立的 DMA 引擎和通信引擎，如果 matmul 的结果留在片上，DMA 可以把数据直接搬到通信缓冲区，Cube 单元同时开始下一轮计算。这种流水线重叠的效果在长序列、大 batch 的场景下尤为明显，因为通信数据量大，流水线的收益也大。
MC2 算子的实现复杂度比单纯融合两个操作高得多。矩阵乘的 tiling 策略要跟通信的缓冲区管理对齐，all-reduce 的分块大小要跟 Cube 单元的输出块大小匹配，不然就得在中途做一次 reshape，又引入了额外的内存操作。ops-transformer 在 mc2/common 目录里维护了一套 tiling 和 buffer 管理的公共逻辑，让各个 MC2 算子共享这些基础设施。
使用前后的效率对比
把 ops-transformer 的算子应用到实际的大模型推理场景中，效率提升主要体现在显存占用和计算吞吐两个维度。下面的对比基于昇腾 NPU 上运行典型 Transformer 推理任务的概括性表现：
| 场景 | 传统方案（框架原生算子） | ops-transformer方案 |
|---|---|---|
| 长序列 attention 显存占用 | 需要存储完整的 n*n 注意力分数矩阵，显存占用随序列长度平方增长 | 分块计算避免存储完整分数矩阵，显存占用接近线性增长 |
| attention 计算延迟 | 大量中间结果在 Global Memory 和片上存储之间来回搬运，带宽成为瓶颈 | 中间结果始终留在片上，带宽消耗大幅降低 |
| MoE token 路由 | 路由打分和 token 重排分散在多个通用算子中，数据布局频繁变换 | 路由流程拆分为专用算子，减少中间数据拷贝和布局转换 |
| MoE 多卡通信 | 矩阵乘和 all-to-all 串行执行，计算和通信之间有闲置等待 | MC2 融合算子让计算和通信流水线重叠，减少闲置等待 |
| FFN+归一化+通信 | 三个操作各自独立，每步都要读写 Global Memory | 融合算子消除中间结果的存储开销，减少内存搬运次数 |
| 整体推理吞吐 | 受限于显存和带宽瓶颈，长序列和大 batch 场景下吞吐受限 | 硬件利用率提升，长序列和大 batch 场景下吞吐显著改善 |
posembedding 与 GMM
除了 attention、MoE、MC2 三大块，ops-transformer 还有一些辅助但同样关键的算子模块。posembedding 目录实现了旋转位置编码（RoPE）相关的算子，包括 kv_rms_norm_rope_cache 这种把 RMSNorm、RoPE 和 KV Cache 操作融合在一起的算子。在大模型推理中，每个 token 进来都要做 RoPE 编码，如果单独调用会引入额外的内存搬运。融合算子把位置编码嵌入到 KV Cache 的更新流程中，省掉了独立的 RoPE 计算和写入步骤。
gmm 目录（Grouped Matrix Multiplication）处理 MoE 场景中按专家分组的矩阵乘法。和普通的 batch matmul 不同，gmm 的每个"组"的大小不一样——每个专家分配到的 token 数量不固定，这导致数据形状不规则。ops-transformer 的 gmm 算子支持这种不规则分组，通过动态 tiling 策略适配不同大小的分组，同时保持 Cube 单元的高利用率。
gmm 的量化版本 grouped_matmul_swiglu_quant_v2 还支持把 SwiGLU 激活函数融合进来，在量化场景下做到矩阵乘和激活函数一次算完。这对 MoE 模型的 FFN 层尤其重要，因为 FFN 层是 MoE 模型里计算量最大的部分。
# gmm + SwiGLU 融合的调用方式（示意）# x 是 permute 后的 token，w_up 和 w_down 是专家权重# 一次调用完成 up_proj -> SwiGLU -> down_projy=ops.grouped_matmul_swiglu_quant(x,w_up,w_down,group_idx)# 量化版本，权重是 int8/int4，计算时在片上反量化gmm 的设计取舍在于：不规则的分组大小意味着无法用固定的 tiling 模板。如果强行对齐到固定大小，会有大量 padding，浪费计算资源。ops-transformer 的方案是根据实际分组大小动态选择 tiling 策略，小分组用 Vector 单元处理（避免 Cube 启动开销），大分组用 Cube 单元处理。这种动态调度增加了控制逻辑的复杂度，但在实际 MoE 工作负载中效果更好。
sparse_flash_attention 与 LightningIndexer
大模型推理领域近年出现了稀疏注意力的方向。核心想法是：大部分 token 的注意力集中在少数关键位置，不需要对整个序列做完整的 attention 计算。ops-transformer 的 sparse_flash_attention 算子支持 KV Cache 的稀疏化存储和注意力计算，配合 lightning_indexer 算子来筛选出哪些 KV 条目值得被保留。
lightning_indexer 的名字来源于一种快速索引策略：它不做完整的 attention 计算，而是用一种低成本的评分方式快速评估每个 KV 条目的重要性，只保留得分最高的那些条目。这在长上下文推理场景下非常有价值——上下文越长，KV Cache 越大，如果不做筛选，显存会被 KV Cache 吃光。
# 稀疏注意力+LightningIndexer 的典型流程# q 是当前 query，kv_cache 是历史 KV 缓存# 先用 indexer 评分筛选，再做稀疏 attentionscores=ops.lightning_indexer(q,kv_cache)# 快速评分keep_mask=ops.topk(scores,ratio=0.3)# 只保留 30% 的重要条目sparse_kv=ops.gather(kv_cache,keep_mask)# 收集重要条目out=ops.sparse_flash_attention(q,sparse_kv)# 只对重要条目做 attention把评分和筛选拆成独立算子是因为两者的数据访问模式差异很大。lightning_indexer 需要扫描整个 KV Cache 但计算量很轻，适合 Vector 单元。sparse_flash_attention 只处理筛选后的子集但计算密集，适合 Cube 单元。分开实现可以让每个算子跑在最合适的硬件路径上。评分算子还有一个反向传播版本 lightning_indexer_grad，支持训练场景下的端到端梯度计算。
总结：
ops-transformer 仓库的目录结构反映了一套清晰的模块化设计。attention 目录专门放 attention 类算子，moe 目录放 MoE 路由算子，mc2 目录放通信融合算子，gmm 目录放分组矩阵乘法，posembedding 目录放位置编码。common 目录放跨模块的公共基础设施，比如 tiling 策略、类型定义、内存管理工具。每个算子子目录下面又分 framework（框架适配层，支持 PyTorch ONNX 等框架的调用）和内核实现层。
这种结构的优势在于：不同模块的开发者可以并行工作，不会互相干扰。公共基础设施集中在 common 里，避免重复实现。框架适配层和内核实现层分离，使得同一个算子可以适配不同框架，而内核代码只需维护一份。
仓库链接：https://atomgit.com/cann/ops-transformer
