一文了解AOL算子加速库

过去一年，随着ChatGPT的发布与快速迭代，基于大数据量、大参数量、大算力的预训练大模型已成为人工智能产业的主要路线。大模型的普及与发展不仅依靠模型本身的创新，更依赖于算力底座的支撑以及软件生态的繁荣，需要伙伴和开发者的积极参与。为加速大模型算力释放，昇腾提供AOL（Ascend Operator Library）算子加速库，不仅覆盖Softmax、MatMul等基础算子，也包括了大模型结构泛化的Flash Attention等高性能融合算子，开发者可以直接使用昇腾内置的算子加速库使能大模型创新与应用。

1 为什么需要算子加速库

算子实际上是面向深度学习任务高度优化的函数，它是网络中层或者节点的计算逻辑，开发者在实现这些计算逻辑的时候，不仅要实现功能逻辑，还要考虑硬件指令的适配，如何支持不同大小/类型的输入，以及如何针对不同的输入Shape进行切分，并保证算子的运行性能。可见，一个算子的代码如此复杂，如果神经网络中的每一个算子都由开发人员实现一遍，那么可想而知，神经网络构建的工作量有多大，难度有多高，特别是面对当前日益复杂的大模型。能否针对某种计算逻辑开发一个通用实现？这样，开发者便像搭房子一样构建网络而无需关注具体的算子实现。

为此，昇腾算子加速库提供了丰富的经过深度优化的高性能算子，且更亲和昇腾AI处理器，可供开发者直接调用，使能大模型极致性能优化，正是这些算子组成了能够让AI任务在昇腾硬件上飞速执行的弹药库。

2 昇腾基础算子加速库

昇腾基础算子加速库包括1400+种高性能算子，其中包括常用深度学习算法计算类型的NN（NeuralNetwork）算子库、能够提供高性能媒体数据处理能力的DVPP（Digital Vision Pre-Processing）算子库、分布式训练中提供卡间高效数据传输的HCCL（Huawei Collective Communication Library）算子库等。

NN（NeuralNetwork）算子库：覆盖了包括TensorFlow、Pytorch、MindSpore、ONNX等框架在内的常用深度学习算法的计算类型，包括Softmax、MatMul、Convolution等典型算子，在算子库中占有最大比重。
BLAS（Basic Linear Algebra Subprograms）算子库：BLAS为基础线性代数程序集，是进行向量和矩阵等基本线性代数操作的数值计算库，CANN支持通用的矩阵乘和基础的Max、Min、Sum、乘加等基础数学运算。
DVPP（Digital Video Pre-Processor）算子库：DVPP是昇腾AI处理器内置的图像处理单元，支持高性能视频编解码、图片编解码、图像裁剪缩放等预处理能力，通过DVPP算子可快速体验媒体处理硬加速能力。
AIPP（AI Pre-Processing）算子库：AIPP可用于在AI Core上完成图像预处理，包括图像尺寸更改、色域转换（转换图像格式）、减均值/乘系数（图像归一化），并与模型推理过程融合，满足推理输入要求。
HCCL（Huawei Collective Communication Library）算子库：HCCL是基于昇腾硬件的高性能集合通信库，提供单机多卡以及多机多卡间的数据并行、模型并行集合通信方案。支持AllReduce、Broadcast、Allgather、ReduceScatter、AlltoAll等通信原语，Ring、Mesh、HD等通信算法，在HCCS、RoCE和PCIe高速链路实现集合通信。

3 昇腾融合算子加速库

在主流大模型网络模型中，大量使用典型的Multi-Head Attention结构，带来了巨大的计算和内存开销。其运行过程中，矩阵乘和softmax结果存放在片上内存会带来巨大的内存消耗，访存性能严重下降，甚至会导致模型无法正常运行，同时网络中的矩阵和向量计算串行执行，也会导致硬件算力发挥受限。

斯坦福的Tri DAO提出了FlashAttention融合算子，其原理是对attention处理过程进行切分和计算等价，使得attention的多个步骤在一个算子中完成，并且通过多重循环、每次处理一小部分数据，以近似流式的方式访问片上内存，减少了片上内存访问的总数据量，并能够将计算和数据搬运更好的重叠隐藏。

CANN针对昇腾AI处理器的片上内存和缓存大小，以及数据搬运通路，优化实现FlashAttention融合算子，充分利用片上缓存，提升Attention处理性能。根据实测，在一些典型场景中，CANN的FlashAttention算子相比融合前的小算子取得了5倍以上的性能提升。开发者可直接调用相关算子接口使能大模型极致性能优化。

Flash Attention算子融合机制

由此可知，算子融合（Operator Fusion）基本操作就是将多个连续的操作或算子合并成一个单一的算子，常用于深度学习和其他计算密集型任务中，以减少计算和内存开销。使用算子融合有助于提高执行效率，减少数据传输和临时存储，以及提高缓存利用率。

除了Flash Attention类算子，昇腾还提供了MoE结构下的高性能融合算子、针对分布式场景提供MC²（Merged Compute and Communication）通算融合类算子等。

MoE （Mixture-of-Experts，混合专家系统）结构将稠密网络的FFN层扩展成具有相同结构的专家网络，通过路由或门控网络决定激活哪些expert进行计算，从而能够实现在低成本下训练更大的模型。昇腾针对MoE结构提供了诸多高性能融合算子，比如GroupedMatMul融合算子可支持专家分组能力、使能多专家多核并行计算；MoeGatingTopKSoftmax将gating部分的softmax以及topk计算融合；MoEInitRouting、MoEFinalizeRouting对routing进行加速处理以提升MoE结构在NPU上的执行性能。

MC²通算融合类算子的工作原理是将通信和计算的操作过程融合，使得计算和通信可以流水并行提升性能，典型算子包括AllGatherMatMul、MatMulReduceScatter、MatMulAllReduce等，这些算子将AllGather、ReduceScatter、AllGather等通信过程和MatMul计算融合，使SP、TP等并行场景下的通信和矩阵计算(Linear)流水并行，可将GPT-3的整网性能提升10%左右。

目前，昇腾已开放部分融合算子代码样例，开发者可以访问https://gitee.com/ascend/cann-ops-adv进行体验和使用。这意味着，开发者不仅可以直接使用昇腾融合算子加速库使能大模型极致性能，也可以从自有的大模型算法和应用出发，基于开放的融合算子代码样例和Ascend C算子编程API开展定制优化，从而提升大模型运行性能或优化资源占用，构建差异化算法竞争力。