前言

Transformer架构自2017年提出以来，已经成为自然语言处理、计算机视觉乃至多模态大模型的基础骨架。随着GPT、LLaMA、DeepSeek等模型参数规模从数十亿跃升至数千亿，模型的推理部署对底层硬件算子库提出了极高的要求。每一层Self-Attention的计算涉及大量矩阵乘法、Softmax归一化和向量运算，这些操作在昇腾NPU上的执行效率，直接决定了端到端推理吞吐量和首Token延迟。

CANN（Compute Architecture for Neural Networks）是昇腾AI处理器的核心计算框架，负责将高层计算图映射到NPU硬件指令。在CANN的算子生态中，ops-transformer项目定位为Transformer类大模型的进阶算子库，专门针对Attention机制、MoE路由计算、FFN前馈网络以及通算融合场景提供高性能算子实现。该仓库目前在AtomGit上持续活跃开发，已累计合并超过6500个Pull Request，覆盖Atlas A2、A3、950PR、950DT以及Kirin X90等多代昇腾产品。

ops-transformer在CANN算子生态中的定位

CANN的算子体系可以分为几个层次。最底层是基础算子库（如Matmul、Transpose、Softmax等单算子），它们通过Ascend C语言在AI Core上实现基本的张量运算。在此基础上，CANN提供了一套标准化的算子注册和调用机制，包括op_host（Host侧形状推导与Tiling）、op_kernel（Device侧Kernel实现）、op_api（aclnn接口封装）和op_graph（图融合规则）四个核心模块。

ops-transformer是在这套机制之上构建的进阶算子库。它不重复实现底层Matmul或Add，而是通过算子融合和通算编排，将多个基础操作组合成一个高效率的复合算子。例如，将Q/K/V投影、注意力分数计算、Softmax归一化、上下文加权求和、输出投影等多个步骤融合为单个Kernel调用，避免中间结果在HBM（高带宽内存）和SRAM之间的反复搬运。

从仓库的目录结构可以清晰看到这一设计。项目顶层分为attention、ffn、gmm（Grouped Matmul）、moe、mc2（MoE Compute Collective Communication）、posembedding、mhc等算子分类目录，每个目录下按统一规范组织op_host、op_kernel、op_api、op_graph子目录。common目录提供全局共享的头文件和工具函数，cmake目录管理编译配置，experimental目录则开放给社区开发者贡献自定义算子。

算子分类体系与核心算子

Attention类算子

Attention是Transformer计算的核心瓶颈。ops-transformer在attention目录下实现了多种变体，覆盖不同模型架构和优化路径：

• flash_attention_score：经典Flash Attention的昇腾NPU实现，通过在线Softmax（Online Softmax）和分块计算减少HBM访问次数，是目前大多数推理框架在昇腾上的默认注意力算子。
• fused_floyd_attention / fused_floyd_attention_grad：针对特定模型架构设计的融合注意力算子，将更多前后操作合并到注意力计算流程中。
• kv_quant_sparse_flash_attention：结合KV缓存量化与稀疏注意力的高阶变体，适用于长上下文场景下控制显存占用。
• sparse_flash_attention：利用注意力矩阵稀疏性跳过部分计算，在特定任务（如文档理解）中能显著减少无效计算。
• mla_preprocess / mla_preprocess_v2：针对Multi-head Latent Attention（MLA）的前处理算子，支持DeepSeek等采用MLA架构的模型。
• attention_update / attention_worker_scheduler / gather_pa_kv_cache：围绕注意力计算的辅助算子，处理KV缓存更新、并行调度和Prefix-Attention场景下的KV Gathering。

从仓库的近期提交记录来看，Attention类算子的优化非常活跃。例如SMLAG（Softmax-L1 Attention Gradient）的提交中，开发者将cmpSoftmaxL1的计算公式改为SimpleSoftmax，并对dSinks进行了优化——每个AIV核只执行一次atomicAdd，同时去除SCFA阶段中oriKv的冗余Gather和Scatter操作。这类修改看似细节，但在高吞吐推理场景下，每次减少一次原子操作就能带来可观的性能提升。

FFN与Gated GLU类算子

FFN（Feed-Forward Network）层在现代Transformer中经历了显著演化。早期的ReLU FFN已被SwiGLU、GeGLU等门控线性单元取代。ops-transformer在ffn目录和experimental目录下提供了相应的融合算子。

近期合并的BiasGateGelu算子是一个典型案例。该算子实现了GeGLU（GELU Gated Linear Unit）的完整计算流程：将输入按列均分为Gate和Value两部分，分别加偏置后对Gate部分施加GELU激活函数，再与Value部分逐元素相乘。如果用基础算子分别调用Add、GELU、Mul，中间结果需要三次HBM写入和读取；而融合算子在SRAM中一次性完成全部计算，写入次数降为一次。

MoE与通算融合算子

MoE（Mixture of Experts）是大模型提升参数效率的关键技术，也是算子设计复杂度最高的场景之一。ops-transformer通过moe、gmm、mc2三个目录覆盖MoE的完整计算链路：

• moe目录：包含moe_compute_expert_tokens等算子，处理Token到Expert的路由分配和Token计数。
• gmm目录：实现Grouped Matmul（分组矩阵乘法），是MoE推理中Expert并行计算的核心算子。仓库中包含grouped_matmul、grouped_matmul_swiglu_quant_v2、grouped_matmul_finalize_routing等多个变体。MegaMoE相关的dispatch_ffn_combine_bf16 Kernel针对Atlas A2（arch22）做了专门的bf16优化。
• mc2目录：提供MoE计算中的集合通信融合算子，如matmul_allto_all、attention_to_ffn、ffn_to_attention。这类算子将矩阵乘法与All-to-All通信操作融合，减少通信与计算的串行等待。

mc2目录中的attention_to_ffn和ffn_to_attention算子尤为关键。在MoE推理中，Attention层输出后需要经过路由分发到不同Expert，Expert计算完成后还需要聚合。如果每一步都独立执行矩阵乘法和通信原语，网络开销会急剧增加。这两个融合算子将相邻阶段的计算与通信绑定执行，实现了流水线化处理。

Host-Device分层架构

ops-transformer的每个算子都遵循统一的四层架构，这一设计贯穿整个仓库。理解这四层的关系，是掌握ops-transformer开发模式的前提。

op_host层：调度与分块策略

op_host层运行在CPU（Host侧），负责算子注册、输入校验、输出形状推导（InferShape）和Tiling（数据分块策略）。Tiling是算子优化的核心环节之一。昇腾NPU的AI Core拥有片上统一缓冲区（UB，Unified Buffer），容量有限但带宽极高。Tiling的任务是将大张量划分为适合UB容量的子块，让Kernel在片上完成计算，避免频繁访问HBM。

每个算子的Tiling实现位于op_host目录下的${o{p}_{n}ame}_{t}iling.cpp文件。仓库目录结构文档明确指出，某些算子针对不同子场景提供独立的Tiling文件，如${op_name}tiling${sub_case}.cpp，这意味着同一个算子会根据输入规模、数据类型等条件选择不同的分块策略。

# cmake/symbol.cmake 中 ophost 的链接库配置片段
# ophost 符号依赖 runtime、acl_rt 以及 mmpa 等基础库
set(OPHOST_LINK_LIBS
    runtime
    acl_rt
    mmpa
)

WHY：op_host编译产物需要链接runtime和acl_rt库才能访问CANN的基础API（内存分配、Stream同步等）。mmpa（Memory Manipulation and Processing Acceleration）提供向量化的内存拷贝和填充操作，在Host侧Tiling数据打包时频繁使用。近期提交中修复mmpa链接缺失的问题，恰恰说明了这些基础库在实际编译中的必要性。

op_kernel层：AI Core Kernel实现

op_kernel层运行在昇腾NPU的AI Core上，使用Ascend C语言编写。这是算子性能的关键所在。每个算子有一个主入口文件${op_name}.cpp，包含GM（Global Memory）到UB的数据搬运、计算逻辑和结果写回。

Kernel实现的核心思路是：按Tiling策略划分数据块，通过DataCopy指令将数据从GM搬入UB，在UB中完成向量或矩阵运算，再将结果搬回GM。对于需要多迭代处理的算子（如Flash Attention需要逐行处理注意力矩阵），会使用循环展开和double buffer策略，在当前块计算的同时预取下一块数据。

op_api层：aclnn接口封装

op_api层提供aclnn（Ascend Computing Library for Neural Networks）标准接口，使上层推理框架（如MindIE、ONNX Runtime）能够以统一的API调用算子。接口封装将op_host的Tiling逻辑和op_kernel的Kernel调用包装为一个函数调用，对用户屏蔽底层细节。

仓库中attention/flash_attention_score/framework子目录下提供了ONNX插件形式的NPUFlashAttention实现。以下展示了框架适配层的一般模式：

# 框架适配层的ONNX插件注册模式（伪代码示意）
# 将CANN aclnn算子注册为ONNX Runtime自定义算子
# 框架层负责将ONNX节点属性映射到aclnn接口参数
# 并管理Host侧Tiling和Device侧Kernel的调用生命周期

WHY：框架适配层的作用是将CANN原生算子暴露给主流推理框架。ONNX插件通过注册自定义节点将CANN算子嵌入ONNX计算图，在运行时直接调用aclnn接口，避免跨框架的序列化开销。这也是仓库在2025年12月新增transformer类ONNX算子插件的技术基础。

类似地，允许ONNX模型直接使用该算子。类似地，common/src/framework下的NPUMultiHeadAttention和moe/moe_compute_expert_tokens/framework下的NPUMoeComputeExpertTokens都是框架适配层的具体实现。

op_graph层：图融合规则

op_graph层运行在CANN图编译器中，负责在计算图优化阶段识别可融合的算子模式（Fusion Pattern），将多个基础算子合并为一个复合算子调用。这是ops-transformer性能优势的重要来源——即使上层框架没有显式调用融合算子，图编译器也能自动识别融合机会。

每个算子的op_graph目录下包含${op_name}_proto.h（算子原型定义，用于模式匹配）和fusion_pass子目录（融合规则实现）。通过定义输入输出的张量属性和算子间的拓扑关系，融合引擎可以在图编译阶段自动完成算子合并。

核心优化策略剖析

算子融合：减少内存搬运次数

算子融合是ops-transformer最核心的优化手段。Transformer推理的计算模式天然适合融合——每一层Attention的计算都是固定顺序的算子链：投影 -> 乘法 -> 加偏置 -> Reshape -> Softmax -> 加权求和 -> 投影。如果每个算子独立执行，中间张量需要在GM和UB之间反复搬运，HBM带宽成为瓶颈。

ops-transformer的做法是将整个计算链封装为单个复合算子。在Kernel内部，中间结果始终驻留在UB中，直到最终结果写回GM。以Flash Attention为例，标准实现需要O(N^2)的HBM空间存储完整注意力矩阵；而融合实现通过在线Softmax技术，只需O(N)的额外空间，且注意力分数无需回写GM即可完成加权求和。

仓库中的mc2类算子（attention_to_ffn、ffn_to_attention、matmul_allto_all）将融合策略推向更高维度——不仅融合计算操作，还融合计算与集合通信。在MoE推理场景中，Expert计算完成后通常需要执行All-to-All通信将结果分发回对应设备。传统方案中矩阵乘法和通信是两个独立的操作，需要等待前一个完成才能启动后一个。融合算子将通信发起嵌入计算流的尾部，实现计算与通信的流水线重叠。

Tiling分块：适配片上存储

昇腾NPU的AI Core片上UB容量通常在数MB级别，而大模型推理中的KV缓存可能达到数百MB甚至数GB。Tiling必须合理划分数据，既不能让单块超出UB容量，又不能分得过细导致Kernel启动开销占比过高。

ops-transformer的Tiling实现位于各算子op_host目录下的tiling文件中。Tiling策略通常会根据输入维度（seq_len、head_dim、num_heads）和NPU型号动态计算分块大小，并考虑double buffer的空间需求——同时容纳当前计算块和下一个预取块。

仓库近期的MXFP8相关提交（fix tiling sink no tilingkey bug）揭示了Tiling实现的一个工程细节：某些算子支持多种Tiling策略（TilingKey），在不同场景下使用不同的分块方式。Tiling下沉（Tiling Sink）优化将Tiling计算从Host侧移到Device侧，减少Host-Device交互延迟，但需要确保TilingKey正确传递，否则会导致运行时错误。

# 仓库编译命令（build.sh）的核心流程
bash build.sh --model_name=<算子名> \
              --soc_version=<NPU型号> \
              --python_path=/usr/bin/python3 \
              --ascend_cann_path=/usr/local/Ascend/ascend-toolkit/latest

WHY：编译单个算子需要指定目标NPU型号（soc_version），因为不同代际的昇腾芯片AI Core微架构存在差异，Tiling策略、UB大小、向量单元宽度等参数不同。build.sh脚本会根据soc_version选择对应的配置文件（位于op_host/config/${soc_version}/目录下），确保编译产物与目标硬件匹配。这也是为什么仓库在op_host/config目录下按NPU型号组织二进制配置和SimplifiedKey信息。

AIV核并行与Kernel流水线

昇腾NPU的AI Core内嵌多个AIV（AI Vector）计算核，可以并行处理不同的数据块。ops-transformer的Kernel实现充分利用了这一特性，将Tiling划分出的数据块分配到多个AIV核上并行计算。

SMLAG算子的优化提交展示了AIV核并行中常见的工程挑战。在该算子的梯度计算阶段，多个AIV核需要向共享的梯度累加器写入结果，原始实现中每个AIV核执行多次atomicAdd，导致原子操作的序列化瓶颈。优化后将每个AIV核的原子操作压缩为一次，大幅降低了同步开销。

此外，仓库中的sub_pipeline_support.yaml工作流配置表明，ops-transformer正在引入子流水线（Sub-Pipeline）支持，允许Kernel内部的不同计算阶段形成流水线。例如在注意力计算中，可以将Q*K乘法、Softmax计算和V加权三个阶段组织为流水线，当前Token的Softmax计算与前一个Token的V加权并行执行。

数据类型与低比特量化支持

大模型推理对显存占用高度敏感，低比特量化是降低显存需求的主要手段。ops-transformer在数据类型支持方面做了大量工作，覆盖FP16、BF16、FP8（MXFP8）、INT8等多种精度。

仓库中的MXFP8相关代码（common目录下的MXFP8修复）和grouped_matmul_swiglu_quant_v2算子体现了低比特量化的两个方向：权重量化和激活量化。MXFP8是华为定义的8位浮点格式，相比标准FP8（E4M3/E5M2）有不同的指数和尾数分配，能够在保持精度的同时提供更好的硬件计算效率。

grouped_matmul_swiglu_quant_v2将Grouped Matmul与SwiGLU激活和量化操作融合，在MoE Expert计算中一次性完成矩阵乘法、激活函数和输出量化，减少中间结果的精度转换开销。

开发者生态与扩展能力

ops-transformer不是一个封闭的黑盒算子库。仓库在2025年10月新增了experimental目录和NpuOpsTransformerExt工程模板，允许社区开发者基于标准框架开发自定义算子并贡献回主线。

NpuOpsTransformerExt模板的特点是支持PyTorch张量操作的无缝集成，提供自动微分能力和GPU/NPU统一接口。开发者可以用熟悉的PyTorch语法编写算子逻辑，模板自动处理昇腾NPU的Kernel生成和Host侧调度。

仓库的CONTRIBUTING.md定义了完整的贡献流程，包括代码规范、测试要求和PR模板。从大量提交记录中的"二级冒烟测试"和"算子泛化测试"等描述可以看出，项目对算子正确性的验证分为多个层次：单元测试验证单算子功能，泛化测试覆盖不同输入形状和数据类型的组合，冒烟测试验证端到端集成。

仓库还支持CANN Simulator仿真工具，允许开发者在没有物理NPU的环境中编译和调试算子代码，降低开发门槛。

使用前vs使用后

在没有ops-transformer算子库的情况下，Transformer推理在昇腾NPU上只能依赖CANN的基础算子（Matmul、Softmax、LayerNorm、Add等单算子）。以一个标准的Self-Attention层为例，需要依次调用十余个基础算子，每次调用的中间结果都要写回HBM再重新读入。在典型的大模型推理配置（batch_size=1、seq_len=2048、head_dim=128、num_heads=32）下，仅Attention层的HBM访问次数就在数百次量级，HBM带宽成为明显的性能瓶颈。

引入ops-transformer的融合算子后，Attention层从十余个独立算子调用变为1到2个复合算子调用。Flash Attention算子通过在线Softmax和分块计算，将注意力矩阵的HBM访问从O(N^2)降到O(N)。对于GPT类自回归推理，prefill阶段的加速通常在2到4倍之间，具体取决于序列长度——序列越长，减少HBM访问带来的收益越明显。

在MoE场景下，差异更加显著。不使用融合算子时，Expert计算、路由分发和All-to-All通信是串行执行的三个阶段，每个阶段之间都有同步等待。使用mc2类融合算子后，计算与通信形成流水线，端到端延迟可降低30%至50%，吞吐量相应提升。Grouped Matmul融合算子将多个Expert的矩阵乘法合并为一个Kernel调用，减少了Kernel启动开销和AIV核的空转时间。

从数据类型角度，基础算子库对FP8或INT8量化的支持有限，往往需要额外的反量化/量化步骤。ops-transformer的低比特融合算子将量化操作嵌入计算流程，在保持精度的同时避免了额外的数据搬运，显存占用可降低约40%至60%。

总结

ops-transformer项目体现了昇腾AI软件栈在Transformer大模型推理领域的系统性优化思路。从算子融合减少HBM访问，到Tiling分块适配片上存储，从AIV核并行提升计算吞吐，到低比特量化降低显存占用，每一层优化都针对昇腾NPU的硬件特性做了针对性设计。仓库的四层架构（op_host、op_kernel、op_api、op_graph）为算子开发提供了清晰的分层规范，而experimental目录和工程模板则为社区扩展留出了空间。

---

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