前言

大模型推理和训练场景中，Transformer 架构占据了绝对主导地位，而 attention、MoE（Mixture of Experts）、FFN（前馈网络）等核心算子的性能直接决定了整个系统的吞吐量和延时。昇腾 CANN 团队开源的 [ops-transformer](还有个值得关注的点：ops-transformer 的开源版本和华为内部版本存在差异。开源版本删除了部分华为自研的 tiling 策略和对特定 shape 的专门优化，这些优化在内部版本里能进一步提升性能。如果你用的是开源版，不要期望拿到跟商用版本一样的性能数字。内部版本通常会针对昇腾硬件做定向的手工汇编优化，开源版本的 tiling 策略虽然泛化能力更强，但在某些特定 shape 上不如手工优化版本。

正是为解决这些算子在昇腾 NPU 上的高效实现而生的进阶算子库。它不只是一个算子集合，更是一套面向 Transformer 类大模型的完整计算加速方案，涵盖了从底层 kernel 编写、 tiling 策略优化，到量化精度保障、再到 ONNX 前端集成的全链路能力。理解并掌握 ops-transformer 的设计思路和用法，是充分发挥昇腾硬件算力的必经之路。

一、为什么需要 ops-transformer

1.1 通用算子库的局限性

在深度学习框架的早期阶段，attention 计算通常被拆解为一系列基础算子的组合：matmul 做 QK 乘法 → softmax 做归一化 → 再一次 matmul 做输出投影。每一步都要读取和写出全局内存，中间的中间结果（intermediate）散落在 HBM 上，反复搬运的开销在序列长度（sequence length）较短时还不明显，但当模型扩展到 7B、70B 参数，序列长度达到 8K、32K 甚至更长时，O(N²) 的注意力计算量和 O(N²) 的内存访问量就成了性能瓶颈。单纯依赖框架自带的通用算子，开发者能做的优化空间极为有限——tiling 策略不可控、融合时机不透明、量化路径需要自己拼接。

1.2 融合算子的必要性

ops-transformer 的核心思路是将 attention 相关的多个基础算子融合为单个原子算子（fused kernel），一次性完成 QK^T、scale、mask、softmax、dropout、乘以 V 的完整计算流程。这样做的好处有两层：第一，显著减少了 HBM 访存次数，热数据可以一直保留在 NPU 的 Unified Buffer 里反复使用；第二，将控制逻辑集中到一个 kernel 里，编译器可以更精准地做指令调度和流水线规划。类似地，MoE 场景下的 token 分发、top-k 门控、加权路由、experts 计算、结果收回等步骤，也被封装为 moeGatingTopK、moeTokenPermute、moeFinalizeRouting 等专用算子，避免了多次跨芯片通信的开销。

1.3 量化加速的支持

大模型推理对显存和带宽极度敏感，FP16/BF16 的全精度存储在超大规模模型上几乎不可承受。ops-transformer 内置了 FP8（E5M2/E4M3FN）、Int8、Int4 等多种量化路径，配合 flashAttentionScore 的 dScaleQ/dScaleK/dScaleV 参数，可以在算子内部完成量化和反量化，开发者无需在 Python 层手动管理量化状态，规避了因精度截断引入的数值误差。

二、整体架构与目录组织

ops-transformer 的源码按照算子功能模块划分，主要子目录包括：

目录	职责	典型算子
attention	自注意力机制算子	FlashAttentionScore, MLA, NSA, RingAttention
moe	MoE 路由与计算	MoeGatingTopK, MoeTokenPermute, MoeFinalizeRouting
mc2	通信算子融合	MatmulAlltoAll, AttentionToFFN, FFNToAttention
ffn	前馈网络融合	GroupedMatMulSwiGLU, GroupedMatMulSwiGLUQuant
gmm	分组矩阵乘法	GroupedMatMul, MegaMoe
mhc	多头缓存相关	MhcPost, MhcPre, MhcRes
posembedding	位置编码	KVRmsNormRoPECache, RoPEMatrix
common	公共框架代码	算子基类、tiling 策略、量化工具
experimental	实验性算子	NpuOpsTransformerExt, MHC 系列
examples	端到端调用示例	FastKernelLaunch, MC2, MoE
tests	单元测试与集成测试	UT/IT 覆盖各算子
docs	开发指南与使用文档	快速入门、调试调优

这种模块化设计的优势在于，每个子目录基本可以独立编译和测试，算子之间的依赖关系清晰，开发者可以按需切入自己关心的模块，而不必面对一个庞大的单体仓库。

三、FlashAttentionScore 算子深度解析

3.1 算子功能与数学公式

FlashAttentionScore 是 ops-transformer 中最核心的算子之一，它实现了 self-attention 的完整前向计算。算子支持两种计算模式，由 pseType 参数控制 add 和 mul 的执行顺序：

当 pseType=1 时，注意力计算遵循公式：

attention_out = Dropout(Softmax(Mask(scale * (pse + (query * d_scale_q) * (key * d_scale_k)^T), atten_mask)), keep_prob) * (value * d_scale_v)

当 pseType≠1 时，公式变为：

attention_out = Dropout(Softmax(Mask(scale * ((query * d_scale_q) * (key * d_scale_k)^T) + pse, atten_mask)), keep_prob) * (value * d_scale_v)

这个细节非常重要——FlashAttention 选择分块（tiling）策略来处理注意力计算，是因为昇腾 NPU 的 Unified Buffer（L1/L2 级缓存）远大于传统 GPU 的 Shared Memory，分块能让热点数据（Q、K、V 的局部块）持续驻留在高速缓存中，避免了对 HBM 的反复读写。算子内部将 S（序列长度）维度切分为多个 tile，每个 tile 内完成完整的 softmax 计算，再通过在线 softmax 技巧（online softmax）将多个 tile 的结果归并为最终输出，这样既保证了数值等价性，又将内存访问复杂度从 O(N²) 降到了 O(N²/B)，B 是 tile 大小。

3.2 核心输入输出参数

// 输入 tensor
query          // [B, S, N, D] 或 [B, N, S, D]，B=batch, S=seq_len, N=heads, D=head_dim
key            // 与 query 形状一致
value          // 与 query 形状一致
realShift      // 可选，位置编码 pse，shape 与 query 的 S 维度对齐
dropMask       // 可选，UINT8，1=保留 0=丢弃
attenMask      // 可选，BOOL/UINT8，1=该位置不参与计算

// FP8 量化参数（可选）
dScaleQ        // query 的全局量化 scale
dScaleK        // key 的全局量化 scale
dScaleV        // value 的全局量化 scale

// 属性
scaleValue     // 缩放系数，默认为 1.0
keepProb       // dropout 保留概率，默认为 1.0
pseType        // 控制 add/mul 顺序，支持 0/1/2/3

// 输出 tensor
softmaxMaxOut  // 中间结果，用于反向计算
softmaxSumOut  // 中间结果，用于反向计算
attentionOut   // 最终输出，shape 与 query 相同

3.3 aclnn 调用示例

#include "acl/acl.h"
#include "acl/ops/acl_fusion.h"

// 创建 context
aclrtContext ctx;
aclrtSetDevice(0, &ctx);

// 输入 tensor 描述
std::vector<int64_t> qShape = {batch, seqLen, numHeads, headDim};
aclTensorDesc *queryDesc = aclCreateTensorDesc(
    ACL_FLOAT16, 4, qShape.data(), ACL_FORMAT_ND);

// 设置 FP8 量化参数
aclScalar *scaleQ = aclCreateScalar(&dScaleQVal, ACL_DT_FLOAT);
aclScalar *scaleK = aclCreateScalar(&dScaleKVal, ACL_DT_FLOAT);
aclScalar *scaleV = aclCreateScalar(&dScaleVVal, ACL_DT_FLOAT);

// 调用算子
aclopExecutor *executor;
aclopCreateExecutor(&executor);
aclopSetOperatorInput(executor, 0, queryDesc, queryDevBuf);
aclopSetOperatorInput(executor, 1, keyDesc, keyDevBuf);
aclopSetOperatorInput(executor, 2, valueDesc, valueDevBuf);
aclopSetOperatorAttr(executor, "scaleValue", scaleValAttr);

// 执行（内部完成 tiling 循环、在线 softmax、dropout mask 生成）
aclrtStream stream;
aclrtCreateStream(&stream);
aclopExecute(executor, stream);
aclrtSynchronize(stream);

为什么这里要单独设置 dScaleQ/K/V？因为在 FP8 场景下，输入数据本身已经是 E5M2 或 E4M3FN 格式的整数量化值，算子需要知道原始数值的真实范围才能正确还原出缩放后的浮点结果。如果不传这些 scale 参数，算子会默认在 BF16/F16 精度下运行，无法利用 FP8 的计算密度优势。

3.4 约束与限制

FlashAttentionScore 在不同硬件平台上的支持情况有差异：

数据类型	Ascend A2 训练	Ascend A3 训练	Ascend 950PR/DT
BF16/F16/FP32	支持	支持	支持
FP8 E5M2/E4M3FN	不支持	不支持	不支持
HIFLOAT8	不支持	不支持	不支持

在 shape 约束上，B 支持 1~2M，N（头数）支持 1~256，S（序列长度）支持 1~1M，D（头维度）支持 1~768。FP8 量化场景下 D 缩减为 1~128，这是因为 FP8 矩阵乘法的最小计算粒度要求更严格。还需要注意，当计算量过大时（比如 B 和 S 同时较大），kernel 执行可能触发 aicore 超时错误，此时建议在模型层面做轴切分，将 B 或 S 维度拆分到多个设备上并行处理。

四、MoE 算子体系：从 Token 分发到结果收回

4.1 MoE 计算的痛点

MoE 架构的核心思想是用一组独立的"专家"（expert）网络替代单一的前馈网络，每个 token 只激活少数几个 expert 做计算。这样做模型容量可以做得很大，但计算模式的复杂性也大幅提升：一个 token 需要先通过门控网络判断自己该去哪些 expert，据此通过通信原语（all-to-all）把 token 分发到对应 expert，计算完成后结果再通过 all-to-all 收回并加权合并。如果这些步骤各自独立实现，token 会在 host 和 device 之间多次搬运，中间结果的临时 buffer 会占用大量显存。ops-transformer 把这套流程拆解为一组协同工作的算子，每个算子专注于一个特定阶段，通过流水线编排来最大化设备利用率。

4.2 核心算子链路

// 第一步：初始化路由上下文，确定 token 的目标 expert
aclTensorDesc *routingDesc = aclCreateTensorDesc(...);
aclopSetOperatorInput(executor, "tokens", tokensDesc, tokensDev);
aclopSetOperatorInput(executor, "routing_weights", rwDesc, rwDev);
aclopExecute("MoeInitRouting", executor, stream);

// 第二步：计算 top-k 门控分数，确定每个 token 激活哪些 expert
aclopExecute("MoeGatingTopK", executor, stream);
// gate 分数决定了 token 的路由路径，选 top-k 而非 top-1
// 是为了让负载更均衡，避免某个 expert 收到过多 token 造成拥堵

// 第三步：分发 token 到对应 expert
aclopExecute("MoeTokenPermute", executor, stream);
// TokenPermute 内部会处理非连续 tensor 的打包和跨设备路由

// 第四步：Expert 计算完成后，收回 token 并恢复原始顺序
aclopExecute("MoeTokenUnpermute", executor, stream);

// 第五步：最终路由合并
aclopExecute("MoeFinalizeRouting", executor, stream);

4.3 为什么需要这么多算子而不是一个 fused kernel

这个问题很自然——为什么不把整个 MoE 流程做成一个巨大的 fused kernel？答案在于通信与计算的解耦。MoE 的 all-to-all 通信需要跨芯片同步，如果把它和计算融合在一起，通信等待时间会阻塞计算单元的流水线。拆分为独立的通信算子和计算算子后，框架可以在这两者之间插入其他计算任务，实现计算与通信的重叠（overlap）。ops-transformer 的算子链路设计，天然支持这种流水线编排。

五、量化实战：FP8 与 Int4 的接入路径

5.1 量化场景下的精度保障

量化推理最大的挑战不是性能，而是精度损失的可控性。FP8 相比 BF16，数值表示范围从 3.4e38 骤降到 57344，精度从 7 位十进制有效数字下降到约 3 位。ops-transformer 在算子层面做了几件事来保证量化精度：在线反量化（on-the-fly dequantization）只在真正参与计算的那一步才把 FP8 数据还原为浮点，中间过程尽量保持高精度；per-tensor 级别的缩放因子（dScaleQ/K/V）使得异常值（outlier）的影响可以被局部放大，而不是污染整个 tensor；softmax 计算强制使用 FP32 累加器，即使 QK 是 FP8，乘累加结果也先以 FP32 存储再做缩放。

5.2 配置量化路径的代码片段

# Python 侧通过 PyTorch 插件接入量化 FlashAttention
import torch
import torch_npu

# QKV 已经过离线量化，变成 INT8 张量
q_int8 = torch.randint(-128, 127, (B, S, N, D), dtype=torch.int8, device='npu')
k_int8 = torch.randint(-128, 127, (B, S, N, D), dtype=torch.int8, device='npu')
v_fp16 = torch.rand(B, S, N, D, dtype=torch.float16, device='npu')

# 量化参数（离线标定得到）
scale_q = torch.tensor([0.01], dtype=torch.float32, device='npu')
scale_k = torch.tensor([0.01], dtype=torch.float32, device='npu')
scale_v = torch.tensor([0.02], dtype=torch.float32, device='npu')

# 调用量化版 FlashAttention
out = torch_npu.npu_fusion_attention(
    q_int8, k_int8, v_fp16,
    scale_q=scale_q, scale_k=scale_k, scale_v=scale_v,
    input_layout="BSND", pse_type=0
)

5.3 Int4 量化路径的特殊处理

MoE 场景下，experts 的权重矩阵通常是最大的显存消耗来源，Int4 量化可以将每个权重的存储空间从 2 字节压缩到 0.5 字节，显存占用直接降为原来的四分之一。ops-transformer 在 GroupedMatMulSwiGLU 系列算子中支持 Int4 量化路径，quantMode 参数控制量化粒度。值得关注的是，Int4 的反量化逻辑必须在 kernel 内部完成，不能依赖 Python 侧做预处理，因为 Int4 数据在内存中是按 byte 对齐的（每 byte 存储 2 个值），手动解包容易引入对齐错误。

六、端到端调用示例：Attention 到 FFN 的融合流水线

ops-transformer 不仅提供单个算子，还提供了组合调用的参考实现。以 attention 到 FFN 的融合场景为例，MC2 目录下的 attention_to_ffn 和 ffn_to_attention 算子支持两个阶段之间的数据直接通过核间通信原语传递，无需落回 HBM：

// Attention 计算完成后，结果直通 FFN 层
aclTensorDesc *attnOutDesc = aclCreateTensorDesc(...);
aclTensorDesc *ffnInputDesc = aclCreateTensorDesc(...);

// 使用 attention_to_ffn 做融合通信
aclopExecutor *a2fExecutor;
aclopCreateExecutor(&a2fExecutor);
aclopSetOperatorInput(a2fExecutor, "attention_output", attnOutDesc, attnDevBuf);
aclopSetOperatorAttr(a2fExecutor, "ep_size", epSizeAttr);
aclopExecute("AttentionToFFN", a2fExecutor, stream);

// FFN 计算结果再返回给 attention 层做残差融合
aclopExecutor *f2aExecutor;
aclopSetOperatorInput(f2aExecutor, "ffn_output", ffnOutDesc, ffnDevBuf);
aclopExecute("FFNToAttention", f2aExecutor, stream);

这样做的收益在于：对于一个深层 Transformer block，attention 输出到 FFN 输入之间的数据搬运，原本需要两次 HBM 读写（attention 写回 → FFN 读出），现在通过核间直连（直接内存访问）可以省去这两次 HBM 访问。对于 batch_size=1、seq_len=8192 的推理场景，这种 fused pipeline 带来的延时改善非常可观。

七、性能对比：使用前 vs 使用后

ops-transformer 提供的融合算子，在典型 Transformer 计算场景下相比逐算子调用的 baseline 有显著优势。以下对比基于 Ascend A2 训练系列实测（数据为相对量级描述，非精确数值）：

维度	逐算子调用（Baseline）	ops-transformer 融合算子	改善说明
HBM 访存次数	多次（QK^T、Softmax、V 乘各自读写）	1 次（单 kernel 完成全部计算）	数据复用率大幅提升
kernel 启动开销	N 个独立 kernel 各自启动	1 个融合 kernel 启动一次	调度开销线性减少
FP8 量化精度损失	Python 侧手动反量化，误差累积	算子内在线反量化，精度可控	数值稳定性更好
MoE Token 分发延时	多次 all-to-all 独立调用	算子流水线 overlap 通信与计算	通信延时被计算掩盖
显存占用	中间结果（softmax_max、softmax_sum）显式分配	融合 kernel 内联管理，临时 buffer 复用	峰值显存降低
代码可维护性	Python 层拼接 N 个算子，版本管理复杂	算子版本对齐，接口稳定	开发和迭代成本降低
ONNX 集成	需手动实现每步对应的 ONNX 算子	提供 NPUFlashAttention 等 ONNX 插件	前端接入成本降低

在实际 benchmark 中，对于 7B 参数模型、batch=1、seq_len=4096 的推理场景，使用 FlashAttentionScore 替代逐算子组合，端到端 throughput 提升明显，延时下降幅度视具体配置而定，核心收益来源于 HBM 访问量的减少和 kernel 调度开销的消除。

八、开发与调试环境

ops-transformer 支持 Docker 环境下的快速部署，项目的 QuickStart 文档详细描述了从源码编译到算子调用的完整流程。开发者可以在本地通过 CANN Simulator 进行算子级别的仿真调试，无需真实硬件即可验证 tiling 策略和数值正确性。项目还提供了丰富的单元测试（UT），覆盖每个算子的边界条件、精度边界和异常输入处理，提交 PR 之前可以通过 scripts/run_ut.sh 快速验证改动没有引入回归。

对于需要深度定制 kernel 的开发者，项目的 common 目录提供了 tiling 策略的抽象基类，算子开发者可以继承这些基类并重写 tiling 逻辑，针对特定的 shape 范围做专门优化。算子开发完成后，通过 CMakeLists.txt 的模板参数化编译，可以一键生成支持不同数据类型的二进制。

九、支持的硬件与版本配套

ops-transformer 的开源算子当前支持 Atlas A2/A3 训练系列产品、Ascend 950PR/Ascend 950DT 以及 Kirin X90 系列。源码分支与 CANN 版本严格对应，master 分支可能存在版本不匹配风险，生产环境请务必使用配套的 tag 分支（建议从 release 仓库确认版本映射关系）。项目保持活跃更新，attention、moe、mc2 等核心模块几乎每天都有 commit 合并，新算子和优化持续进入主干。

---

仓库地址：https://atomgit.com/cann/ops-transformer