CANN/ops-transformer FFA算子设计

1 计算过程【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer按照FusedFloydAttention正向计算流程实现整体计算流程如下query, key_1通过特定矩阵乘计算得到a_0与query, key_2通过特定矩阵乘计算得到的a_1相加然后再乘以缩放系数scale_value。此时的结果通过atten_mask进行select操作将atten_mask中为true的位置进行遮蔽得到结果masked_attention_score即atten_mask中为true的位置在select后结果为负的极小值经过softmax计算之后变成0从而达到遮蔽效果。借鉴FlashAttention加速算法FusedFloydAttention使用FlashSoftmax操作对masked_attention_score进行运算用以代替原公式中的softmax运算而后将结果与value做matmul运算。由于FlashSoftmax操作对masked_attention_score的Skv(输入key、value的sequence length)方向进行了切分故实现过程中存在一个刷新流程具体如下每次FlashSoftmax计算只对切分后的一个SkvSplitSkvSplit是针对Skv轴进行切分之后的序列长度的简称进行操作并从第二次循环开始记录exp其中 i 表示Skv切分后的循环变量针对exp的i是从1开始 exp的计算公式如下 $$ exp[i] e^{max_{i - 1} - max_{i}} $$当i 0时计算出的MM[PV]结果直接保存到ub_attention_out[0]的ub中。从i 1开始需要增加Mul和Add操作即将上一次的MM[PV]的结果和当前exp相乘相乘完的结果和本次MM[PV]的结果相加得到的结果保存到ub_attention_out[1]的ub中。以此类推遍历Skv计算完成。由于FlashSoftmax计算中的除sum被后移到输出attention_out之前因此最后需要将ub中的ub_attention_out按行除以softmax_sum并将最终完整的结果保存到输出内存attention_out(Final)上。2 Tiling设计Tiling操作的目的是为了找到一种更高效的NPU执行方式原始的数据量一般是非常大的没有办法通过一次指令调用就完成所有计算因此需要将数据量分到多个核上并行计算且每个核上也需要考虑如何循环计算性能最优不同的输入可能有不同的最优执行方式所以需要通过tiling策略决定怎么将数据分配到各个核上进行计算。根据硬件架构特征AI Core分成AIC和AIV两个独立的核AIC和AIV核拥有自己独立的Scalar计算单元能够独立加载自己的代码段单独执行。AIC和AIV分离的架构可以使得AIC和AIV并行执行。AIC和AIV之间数据交互的通路是L2和GMGlobal Memory高带宽存储器两者之间的交互次数对性能影响是比较大的同时由于AIC和AIV算力差异两者需要使用不同的基本块大小本着尽量减少AIC和AIV通信次数和发挥最大算力的原则CVtiling分离策略应运而生可以有效地减少CV通信次数同时根据不同单元的buffer特征选择不同的基本块进行计算从而提升算子性能。对于FFA算子Vector计算涉及多个输入、输出、中间计算结果、double-buffer设计等需要将buffer分配成多份最优分配方案中最大一份为32KB由于Vector计算使用的数据类型是float32因此Vector的tiling基本块为8 * 1024。为了充分发挥Cube的算力在CV之间一轮计算的数据量进行了4:1的配比又由于Cube侧的输入数据类型是float16输出是float32Cube的基本块为32 * 32所以通过nRatio32配比出32 * 1024的数据量。伪代码如下// C-Tiling: (S1_c_i,D)x(D,S2_c_i) (S1_c_i, S2_c_i):(32,1024) // V-Tiling: (S1_v_i, S2_v_i) (8,1024) // C侧 matmul计算 Bmm((S1_c_i,D)x(D,S2_c_i)) 32*1024 // 输出结果32*1024放到workspace上 // V侧 Vector计算 for S1_c_i/S1_v_i32/8: copy_gm_to_ub(S1_v_i*S2_v_i) // 从bmm的workspace上拷入bmm结果数据 Vector(S1_v_i,S2_v_i) // 进行Vector计算 copy_ub_to_gm(S1_v_i*S2_v_i) // Vector计算结束得到最终输出数据拷贝到GM上 // 由于Cube侧计算数据比Vector侧大因此ub内需要再次进行Vector Tiling从而产生了S1方向的配比S1_c_i/S1_v_i上述示例中仅在S1方向开了配比S2方向C/V计算的长度是一致的当然也可以在S1/S2方向均开启配比这样做的好处是Cube一次可以发射大块的数据避免因为小块数据不断发射带来的通信开销也能最大程度地使用Cube单元的buffer。3 流水设计为了追求极致性能必须充分利用硬件资源通常需要进行不同pipeline的流水设计。流水设计的宗旨是尽量使某一条pipeline达成bound效果使硬件的某一个单元一直在工作达到性能上限。3.1 V侧流水V侧流水设计需要考虑Vector内部的搬运及计算过程实施的优化手段主要是double buffer。 Vec的功能被拆分成2个流水任务Stage1、Stage2每个任务专注于完成单一功能需要处理的数据被切分成2片使用ping-pong表示两个数据处理任务每个任务需要依次搬运DataCopy与计算Clc操作。任务间的箭头表达数据间的依赖关系比如Stage1处理完DataCopy之后Stage2才能对Clc进行处理。 不进行流水设计时搬运与计算任务之间是串行执行的会出现断流现象即第一次DataCopy完成之后的搬运流水就一直处于空闲状态直到第一次搬入的数据计算完成并搬出之后搬运流水才会继续工作进行第二次DataCopyVector计算和搬出流水也存在同样问题。通常这种情况下性能是极差的。对ping-pong流水间的double buffer处理后对于同一片数据搬运DataCopy与计算Clc之间的处理具有依赖关系需要串行处理不同的数据切片同一时间点可以有多个任务在并行处理由此达到任务并行、提升性能的目的。其中ping、pong两块计算数据所占用的内存资源均相互独立。融合算子V侧计算过程较多情况也比较复杂通常简单的double buffer是无法覆盖所有情况的因此会出现不同的计算流水排布。3.2 CV流水对于融合算子V侧的计算是依赖C侧的计算结果的如果只关注V侧流水不关注C侧则C侧与V侧很有可能是串行流水的效果不能达到并行计算的目的无法使得融合算子性能达到最优从而有了CV流水设计。此外CV流水在不同算子情况下表现的现象也是不一致的Cube双发机制可实现两种场景下的流水优化C侧总耗时 V侧总耗时该场景流水特征下Vector计算节点少计算速度快在Atlas A2 训练系列产品 C:V1:2的情况下Cube的搬运时长足以掩盖Vector的计算时长因此只要关注Cube的MTE2耗时即可最终达成MTE2 bound。在Cube双发机制下提前发射两块Cube计算Cube1、Cube2计算可以衔接使得Cube利用率最高达成Cube bound。C侧总耗时 V侧总耗时C侧连续发射两块Cube计算这样可以保证V侧计算完上一轮时可以立马启动当前轮的计算而不用等待Cube1的数据。这样可以使V侧一直在工作达成Vector bound。4 优化效果单算子推理耗时输入shapeB, H, N, M, K, DChunksAscendCs加速比 (Speedup)(1, 32, 256, 256, 256, 32)0.02820.02091.35x(1, 32, 384, 384, 384, 32)0.09550.06041.58x(1, 32, 512, 512, 512, 32)0.19410.13831.40x(1, 32, 768, 768, 768, 32)0.76400.45631.67x(1, 16, 256, 256, 256, 64)0.01580.01421.11x(1, 16, 384, 384, 384, 64)0.05610.03161.78x(1, 16, 512, 512, 512, 64)0.12050.11511.05x(1, 16, 768, 768, 768, 64)0.42490.33021.29x显存占用输入shapeB, H, N, M, K, DChunkGBAscendCGB显存节省率 (Mem Saving)(1, 32, 256, 256, 256, 32)6.382.6558.46%(1, 32, 384, 384, 384, 32)19.975.5172.41%(1, 32, 512, 512, 512, 32)45.509.5079.12%(1, 16, 256, 256, 256, 64)3.882.4038.14%(1, 16, 384, 384, 384, 64)11.535.1855.07%(1, 16, 512, 512, 512, 64)25.509.1364.20%5 附录——FusedFloydAttention算子中IterateBatch配置本方案中涉及两种类型的批量矩阵乘数据格式ND第一种是: NMD NKD, 其中N是batch轴 D是K轴;第二种是: NMD KMD, 其中M是batch轴 D是K轴第一种方案中batch轴在最外轴M和K可能有切片这样每次单batch的ND矩阵是连续的但是不同batch的ND矩阵间可能会存在间隔。比如M256, 假设切片为N[0:16]M[128:256]D[:], 每个batch为 128*D 的数据 batch数据起始地址之间间隔256 * D;针对这种情形我们使用IterateBatch的NORMAL数据排布对于batch数据之间的地址间隔需要配置matrixStrideA matrixStrideB等参数这两个参数能指定输入A矩阵和B矩阵的相邻batch的ND矩阵间的间隔。上述的例子中将matrixStrideA设置为 256*D。第二种方案中batch轴在中间这种情形下每次单batch的ND是不连续的跳跃的。一种方法是实现中自己控制数据搬运使用低阶接口或者使用AscendC中的高阶接IterateBatch。使用IterateBatch针对batch在中间轴的情况可配置使用BSNGD的数据排布格式。其中每个batch计算一个SD数据。例如NMD KMD可以把B矩阵KMD理解为[B1,SK,N1,GM, DD]这种排布这样就能使用IterateBatch高阶接口。使用IterateBatch每次调用时的batch数据需要指定但是不能任意指定有硬件约束。主要是要求指定的计算batch数据能完全搬运进L1 cache(对于910B为512KB)即计算数据量不能超过NPU L1 Cache。下面是相关的计算流程。计算一个批次Batch矩阵乘计算的数据量 (Bytes)oneBatchBytes datatype Bytes * D * (K1 K2)K1与K2取实际计算的矩阵的轴的切片长度比如上面例子中M[128:256]对应的是128实际最大的batch数量则如下maxBatchNums Floor(L1Size / oneBatchBytes) # 向下取整对应的循环次数如下Loops Ceil(TotalBatches, maxBatchNums) # 向上取整【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考