1 计算过程【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer图1 训练计算流程图按照FlashAttention反向计算流程实现整体计算流程如下重计算p$p SimpledSoftmax(Mask(Matmul(query, key^T) pse_shift) * scale)$本步骤重计算了fa流程中的softmax结果p计算结果保存的ub中。计算dp$dp Dropout(Matmul(dy, value^T))$该计算包含matmul计算和dropout计算matmul计算中左矩阵为dy右矩阵为转置后的value。计算ds$ds p * Sub(dp, FlashSoftmaxGrad(dy, attention_in))$本计算中FlashSoftmaxGrad计算的入参为dy、正向输出attention_in该结果与dp做减操作最终的结果与p相乘得到结果ds。计算dq$dq Matmul(ds, key) * scale$本计算将ds结果与key做matmul计算并将结果与scale相乘得到结果dq。计算dk$dk Matmul(ds^T, query) * scale$本计算将转置后的ds结果与query做matmul计算并将结果与scale相乘得到结果dk。计算dv$dv Matmul(DropOut(p)^T, dy)$本计算p的结果做drop计算转置后与dy做matmul计算。2 每个计算阶段的入口以sameAB模板为例得到每个核处理的任务块数后调用函数ProcessCube1阶段流程中1计算入口ComputeMM1Vector1阶段流程中2-3计算入口ComputeVecCube2阶段流程中4-6计算入口ComputeMMDqkv3 Tiling设计Tiling操作的目的是为了找到一种更高效的NPU执行方式原始的数据量一般是非常大的没有办法通过一次指令调用就完成所有计算因此需要将数据量分到多个核上并行计算且每个核上也需要考虑如何循环计算性能最优不同的输入可能有不同的最优执行方式所以需要通过tiling策略决定怎么将数据分配到各个核上进行计算。根据硬件架构特征AI Core分成AIC和AIV两个独立的核AIC和AIV核拥有自己独立的Scalar计算单元能够独立加载自己的代码段单独执行。3.1 Atlas A2 训练系列产品Atlas A2 训练系列产品 芯片AIC和AIV分离的架构可以使得AIC和AIV并行执行。AIC和AIV之间数据交互的通路是L2和GMGlobal Memory高带宽存储器两者之间的交互次数对性能影响是比较大的同时由于AIC和AIV算力差异两者需要使用不同的基本块大小本着尽量减少AIC和AIV通信次数和发挥最大算力的原则CVtiling分离策略应运而生可以有效地减少CV通信次数同时根据不同单元的buffer特征选择不同的基本块进行计算从而提升算子性能。对于FAG算子Vector计算涉及多个输入、输出、中间计算结果、double-buffer设计等需要将buffer分配成多份最优分配方案中最大一份为32KB由于Vector计算使用的数据类型是float32因此Vector的tiling基本块为8 * 1024。为了充分发挥Cube的算力在CV之间一轮计算的数据量进行了1:16的配比又由于Cube侧的输入数据类型是float16输出是float32Cube的基本块为128 * 128所以通过nRatio8配比出128 * 1024的数据量。伪代码如下// C-Tiling: (S1_c_i,D)x(D,S2_c_i) (S1_c_i, S2_c_i):(128,1024) // V-Tiling: (S1_v_i, S2_v_i) (8,1024) // C侧 matmul计算 Bmm((S1_c_i,D)x(D,S2_c_i)) 128*1024 // 输出结果128*1024放到workspace上 // V侧 Vector计算 for S1_c_i/S1_v_i128/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.2Ascend 950PR/Ascend 950DTAscend 950PR/Ascend 950DT同样是AIC和AIV分离的架构保留了AIC AIV并行执行特性同时新增了AIC和AIV之间的高速数据交互通路L0C-UB和UB-L1。降低了CV之间交互的成本流水并行度更高且相比于 Atlas A2 训练系列产品 复杂的tiling切块策略在950上仅使用一种基本块就可以获得较好的性能。对于FAG算子Vector计算涉及多个输入、输出、中间计算结果、double-buffer设计等需要将buffer分配成多份最优分配方案中最大一份为32KB由于Vector计算使用的数据类型是float32因此Vector的tiling基本块为64 * 128由于Cube与Vector核数为12的数量比为了充分利用cube核的算力Cube侧考虑采用128 * 128的基本块即每个cube核计算完128 * 128的数据后均分给两个vector核处理。伪代码如下// C-Tiling: (S1_i,D)x(D,S2_i) (S1_i, S2_i):(128,128) // V-Tiling: (S1_i / 2, S2_i) (64,128) // C侧 matmul计算 Bmm((S1_i,D)x(D,S2_i)) 128*128 // valuedy,qk输出结果128*128放到L0C上 fixp_l0c_to_ub(S1_i / 2,S2_i); v0: 64*128, v1: 64*128 // 通过CV通路L0C-UB均分写到两个V核的UB上 // V侧 Vector计算 Vector(S1_i / 2,S2_i) // 进行Vector计算 copy_ub_to_l1(S1_i/2*S2_i) // Vector计算结束得到输出数据通过CV通路UB-L1两个v核数据聚合拷贝到L1上 Bmm((S1_i,S2_i)*(S2_i,D)); // dq最终结果输出到workspace Bmm((S2_i,S1_i)*(S1_i,D)); // dkv最终结果输出到workspace上述示例中通过两个拷贝指令fixp_l0c_to_ub和copy_ub_to_l1分别将cube侧计算的数据和vector侧计算的数据copy到vector侧和cube侧这两个拷贝都是核内高速通道。4 流水设计为了追求极致性能必须充分利用硬件资源通常需要进行不同pipeline的流水设计。流水设计的宗旨是尽量使某一条pipeline达成bound效果使硬件的某一个单元一直在工作达到性能上限。4.1 V侧流水V侧流水设计需要考虑Vector的搬运及计算过程实施的优化手段主要是double buffer。 以下面的流水任务示意图为例Vec的功能被拆分成2个流水任务subA、subB每个任务专注于完成单一功能需要处理的数据被切分成2片使用ping-pong表示两个数据处理任务每个任务需要依次搬运DataCopy与计算ClcClc表示Vector计算操作。任务间的箭头表示数据间的依赖关系比如subA处理完DataCopy之后subB才能对Clc进行处理。 从图上可以看出不进行流水设计时搬运与计算任务之间是串行执行的会出现断流现象即第一次DataCopy完成之后的搬运流水就一直处于空闲状态直到第一次搬入的数据计算完成并搬出之后搬运流水才会继续工作进行第二次DataCopyVector计算和搬出流水也存在同样问题。通常这种情况下性能是极差的。将上图的流水任务做ping-pong流水间的double buffer处理后流水任务运行起来的示意图如下从运行图中可以看出对于同一片数据搬运DataCopy与计算Clc之间的处理具有依赖关系需要串行处理不同的数据切片同一时间点可以有多个任务在并行处理由此达到任务并行、提升性能的目的。其中ping、pong两块计算数据所占用的内存资源均相互独立。 FAG融合算子V侧计算过程较多情况也比较复杂通常简单的double buffer是无法覆盖所有情况的因此会出现不同的计算流水排布。不同的计算流水适用于不同类的shape特征以达到在该类特征下最好的流水设计。4.2 CV流水融合算子通常包含了Vector计算和Cube计算对于FA算子V侧的计算是依赖C侧的计算结果的如果只关注V侧流水不关注C侧则C侧与V侧很有可能是串行流水的效果不能达到并行计算的目的无法使得融合算子性能达到最优从而有了CV流水设计。此外CV流水在不同算子情况下表现的现象也是不一致的FA/FAG的Cube双发机制又称为CV间preload流水可实现两种场景下的流水优化C侧总耗时 V侧总耗时该场景流水特征下Vector计算节点少计算速度快在 Atlas A2 训练系列产品 C:V1:2的情况下Cube的搬运时长足以覆盖Vector的计算时长因此只要关注Cube的MTE2耗时即可最终达成MTE2 bound。在Cube双发机制下提前发射两块Cube计算Cube1、Cube2计算可以衔接使得Cube利用率最高达成Cube bound。C侧总耗时 V侧总耗时该场景流水特征下Vector计算节点多Vector计算是瓶颈C侧的搬运不足以覆盖V侧的流水因此需要进行CV流水排布尽量达到CV并行的效果最通用的优化手段是C侧提前发射。C侧连续发射两块Cube计算这样可以保证V侧计算完上一轮时可以立马启动当前轮的计算而不用等待Cube1的数据。这样可以使V侧一直在工作达成Vector bound。5 多模板设计为了使不同的输入可以复用相同的tiling和流水采用了模板的方式来实现融合算子但是不同的输入全部使用同一套模板时又无法达到性能最优和功能泛化因此需要根据输入shape的特征区分不同的模板来实现。 FAGFlashAttentionScoreGrad简称FAG融合算子的多模板设计思路主要为根据核内及核间切分进行模板拆分由于硬件buffer大小是有限的而计算的数据量又是巨大的无法一次计算完那么就需要进行tiling切分shape不同会导致算子的切分轴不同而算子的切分轴会影响模板的功能及性能。简单的elewise类算子往往会将所有的轴fuse成一根轴进行切分逻辑简单因此模板也比较单一。而融合算子融合了elewise、broadcast、reduce及matmul等多类场景功能复杂为达到较高的性能要求往往需要根据切分轴进行模板拆分模板拆分时为了达到性能最优需要考虑如下几个点a. 将核心的数量用满防止部分核闲置b. 每一个核心被分配的计算量相对均匀避免出现某些核计算的数据量过大其余核在围观的情况。c. AIC和AIV之间处理的数据量要符合其对应的算力避免AIC或AIV出现长时间的空闲。FAG算子包含B、N2(key和value的N)、G(query_N/kv_N)、S1(query的S)、S2(key和value的S)共5个轴切分顺序是先核内再核间核内切分依据基本块大小选择切分轴核间切分是把核内切分后剩余的轴合并后依据AI Core核数再进行切分。由于shape的大小不同切分轴会发生变化从Vector视角FAG算子划分为如下几类模板模板按序号排优先级序号越小优先级越高越先匹配。FAG算子根据适用范围不同划分成以下几类模板Atlas A2 训练系列产品模板切分轴走入条件适用范围B模板多核切B非FP32bestBasicBlockNum S1 4 ? 64 * 128 / 4 * 3 : 64 * 128; N1 * G * alignedS1 * alignedS2 bestBasicBlockNum。普通场景N2模板多核切N2非FP32S1 * S2 64 * 128。普通场景SameAB模板UB切S1S2多核切S2确定性计算场景非FP32或者S1 1024 or S2 1024走SameAB模板其余走1.2模板。TND场景平均S1 1024 and 平均S2 1024走SameAB模板其余走1.1模板。普通场景S1 1024 or S2 1024。确定性计算场景TND场景普通场景S1S2模板UB切S1S2多核切N2确定性计算场景不满足SameAB模板条件。普通场景不满足其他模板条件。确定性计算场景普通场景TND模板UB切S1S2多核切S2不满足SameAB模板条件。TND场景Ascend 950PR/Ascend 950DT模板切分轴走入条件适用范围BN2多核切BN2S1128 and S2 128 and G1非FP32。普通场景BN2S2模板多核切BN2S2非FP32TNDB * N2 * S2 cubeCoreNum and G 1。普通场景确定性计算模板多核切BN2GS1S2开启确定性计算确定性计算场景BN2GS1S2模板多核切BN2GS1S2上述场景都不满足普通场景每一类模板都有其独特的UB及Block切分轴能处理某一类具备特定shape特征输入的场景针对该类shape特征进行模板设计。根据特殊场景及特定优化进行模板特化每一个融合算子有一个基础模板并辅以多个特化模板基础模板覆盖功能以及大部分此类shape特征输入的性能。特化模板是覆盖特定场景的极致性能主要根据某些适用于特定场景的特殊手段进行的优化不适合进行泛化。例如根据特殊场景空tensor特化而出的empty_input模板基于角色的Cube核管理RCM优化方案设计的确定性计算模板等。根据不同计算流水进行模板特化为了充分发挥硬件优势通常融合算子都需要进行流水设计以提高融合算子性能不同的流水设计对代码的架构影响非常大为了提升代码的可维可测可读性需要根据不同的计算流水进行模板特化达到特定场景的极致性能。5.1 多模板详细设计基本概念CV基本块:表示Cube或者Vector单次计算的数据量大小用来描述一次完整的Cube和Vector交互的数据量通常也等价于核间基本块。在芯片上由于Cube和Vector之间通信有一定开销所以CV基本块设置的比较大一般合适的数据量大小是512*1024单位Bytes。又由于Cube和Vector的核内Buffer有限所以核间基本块可能要通过多次核内计算完成。核内基本块如果CV基本块过大Cube和Vector核内会将CV基本块进一步的切分切分成适合核内L0A、L0B、L0C、UB等大小的基本块这个就叫核内基本块对于Cube侧核内基本块一般在32KB单位Bytes这样可以让L0A、L0B的DoubleBuffer能力展开同时算力和带宽也能尽可能用满在Vector侧一般基本块大小是32KB单位Bytes。xxx.i:表示经过切分后CV基本块中的某根轴的大小, 一般基本块都是2维的xxx.i 表示其中一个维度的大小。xxx可以是B、N2、G、S1、S2; 例如S1 512, S2 1024 S1.i 64, S2.i 128, 表示把[S1, S2]切分成大小是[64, 128]的基本块S1轴的基本块大小是64S2轴的基本块大小是128。xxx.o: 表示经过基本块切分后某根轴的分数xxx可以是B、N2、G、S1、S2;例如S1 512, S2 1024 S1.i 64, S2.i 128, 表示把[S1, S2]切分成大小是[64, 128]的基本块S1.o 512 / 64 8, S2.o 1024 / 128 8 一共切分成8 * 8个基本块。Atlas A2 训练系列产品FAG算子的模板划分如下以下模板序号越大模板的优先级越高序号1的模板是泛化模板支持所有shape核间切分B、N2、G、S1轴核内切分S1轴、S2轴模板tiling代码文件ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_s1s2_bn2gs1s2.cpptiling代码类FlashAttentionScoreGradTilingS1s2Bn2gs1s2kernel代码ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2.h条件支持所有shape其他模板如果不支持就会走到这个模板 依据这个模板按照最通用的做法可以支持所有的Shape。但是由于核内一次只能处理S1 * S2的大小当S1和S2都比较小的时候会有频繁的CV交互开销性能较差。当S1和S2都比较小的时候会路由到下面的这些模板。核间切分B、N2轴核内切分G、S1、S2该模板是通过单纯的分核改变来优化S1、S2都比较小且(B * N2比较大或者G 1)场景下的性能tiling代码文件ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_s1s2_bn2.cpptiling代码类FlashAttentionScoreGradTilingS1s2Bn2kernel代码ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2.h条件(S1 1024) (S2 1024) (B * N2 * 2 CoreNum || G 1)依据S1和S2都比较小且B和N2比较大的时候这时候把B和N2用于分核核内不切分N2.i循环N2.i次进行Cube和Vector计算。核间切分B、N2.o轴核内切分N2.i、G、S1、S2轴, 该模板是为了优化G * S1 * S2都比较小的场景时的性能把N2轴切分到核内并且在核内也切分N2.i轴用于加速Vector计算。相比于模板2, 模板3会更复杂一些模板3在核内计算中也切分了N2.i轴让每次的计算量更大。tiling代码文件ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_ngs1s2_bn.cpptiling代码类FlashAttentionScoreGradUngs1s2BbnTilingkernel代码FlashAttentionScoreGradUngs1s2Bbn条件G * S1 * S2 32KB G 1 S2 1536依据当G * S1 * S2小于32KB时可以通过把N2轴切分一部分到核内让CV基本块更大同时在Vector核内把N2.i的也进行切分让单次Vector的计算量更大提升Vector利用率。核间切分B轴核内计算B.i 、N2、 G、S1、S2轴该模板是为了优化N2 * G * S1 * S2比较小时的性能tiling代码文件ops-transformer-dev/attention/flash_attention_score_grad/op_host/flash_attention_score_grad_tiling_bngs1s2_b.cpptiling代码类FlashAttentionScoreGradUbngs1s2BbTilingkernel代码FlashAttentionScoreGradUngs1s2Bbn条件N2 * G * S1 * S2 64 * 128依据如果希望单纯的把B.i放入CV基本块中那么内层轴N2 * G * S1 * S2就需要足够小一般是根据这个只小于64KB的话Bmm1和Bmm2的数据量一般不会超过L1的一半那么B轴切分时有意义的否则单个Matmul就把L1用满多个Matmul之间的数据搬入没有办法和计算并行。Ascend 950PR/Ascend 950DTFAG算子的模板划分如下以下模板序号越大模板的优先级越高序号1的模板是泛化模板支持所有shape虽然存在多个模板但在实现时仅存在两个模板文件一个是确定性计算另一个是非确定性计算模板其中非确定性计算模板包含了BN2BN2S2BN2GS1S2三种切分模板在代码中通过模板参数隔离各自的实现逻辑核间切分B、N2、G、S1、S2轴模板tiling代码文件ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpptiling代码类FlashAttentionScoreGradTilingUs1s2Bs2Regbasekernel代码ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h条件支持所有shape其他模板如果不支持就会走到这个模板 依据这个模板按照最通用的做法可以支持所有的Shape。但是在部分场景下性能不是最优此时根据不同场景判断路由到以下其他特化模板。核间切分B、N2轴模板tiling代码文件ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpptiling代码类FlashAttentionScoreGradTilingUs1s2Bs2Regbasekernel代码ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h条件S1128 and S2 128 and G1非FP32。 依据S1128 and S2 128 and G1时不存在多核切分S1S2轴即不存在多核累加可以省去前置GM清零和后置cast和musl计算性能更优。核间切分B、N2、S2轴模板tiling代码文件ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpptiling代码类FlashAttentionScoreGradTilingUs1s2Bs2Regbasekernel代码ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel.h条件非FP32TNDB * N2 * S2 cubeCoreNum and G 1。 依据B * N2 * S2 cubeCoreNum and G 1场景dk和dv的计算无需多核累加可以省去前置dk dv GM清零和后置cast和musl计算性能更优。确定性计算模板tiling代码文件ops-transformer-dev/attention/common/op_host/arch-310/flash_attention_score_grad_tiling_s1s2_bn2gs1s2_regbase.cpptiling代码类FlashAttentionScoreGradTilingUs1s2Bs2Regbasekernel代码ops-transformer-dev/attention/common/op_kernel/arch-310/flash_attention_score_grad_kernel_deter.h条件开启确定性计算。 依据确定性计算场景对分核有比较严格的要求通过特定分核方式避免多核同地址累加达到确定性计算的效果。6 编程视角6.1 AscendC高阶APIAtlas A2 训练系列产品当前AscendC高阶API提供了两种编程模式第一种是以Vector为主核Cube为从核的视角Vector0和Vector1会独立发起Matmul的任务两者没有关联性。第二种是以Cube为主核Vector为从核这时会由V0统一发起Matmul任务这个任务的结果由V0和V1共同处理一般是V0、V1各处理一半。当前如果某个模板的结尾是_sab说明这个模板是一个以Cube为主核的模板。例如ops-transformer-dev/attention/flash_attention_score/op_kernel/flash_attention_score_s1s2_bn2gs1_sab.h ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_sab.h以Cube为主核对于FlashAttention来说由于V0、V1的Matmul任务可以复用左矩阵且输出的部分结果可以在L0C累加减少了对于带宽的依赖诉求大部分场景性能会更优。Ascend 950PR/Ascend 950DT为了实现极致性能除部分确定性计算场景全部切换到AscendC低阶API实现且无论是高阶还是低阶都是以Cube为主核实现。6.2 AscendC低阶APIAtlas A2 训练系列产品当前还存在一些没有使用AscendC高阶API的模板例如ops-transformer-dev/attention/flash_attention_score_grad/op_kernel/flash_attention_score_grad_s1s2_bn2gs1s2_basic.h这个模板更加彻底地使用了以Cube为主核Vector为从核这时Matmul的任务都已经完全从Cube侧发起通过同步通知Vector侧。Ascend 950PR/Ascend 950DT除部分确定性计算场景其余场景全部采用低阶API实现可以实现极致的内存复用。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考