SparseFlashAttentionGrad【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明算子功能根据topkIndices对key和value选取大小为selectedBlockSize的数据重排接着进行训练场景下计算注意力的反向输出。计算公式根据传入的topkIndice对keyIn和value选取数量为selectedBlockCount个大小为selectedBlockSize的数据重排公式如下$$ selectedKey\text{ }\text{ }Gather \left( key,topkIndices \left[ i \left] \left) ,\text{ }0\text{ } i \text{ }selectBlockCount\right. \right. \right. \right. $$$$ selectedValue\text{ }\text{ }Gather \left( value,topkIndices \left[ i \left] \left) ,\text{ }0\text{ } i \text{ }selectBlockCount\right. \right. \right. \right. $$阶段1根据矩阵乘法导数规则计算$dP$和$dV$:$$ dP\mathop{{}}\nolimits_{{t,:}}dO\mathop{{}}\nolimits_{{t,:}}\text{}V\mathop{{}}\nolimits^{{T}} $$$$ dV \left[ u \left] P\mathop{{}}\nolimits_{{T}}^{{t,:}}\text{}dO\mathop{{}}\nolimits_{{t,:}}\right. \right. $$阶段2计算$dS$:$$ d\mathop{{S}}\nolimits_{{t,:}} \left[ P\mathop{{}}\nolimits_{{t,:}} \left( dP\mathop{{}}\nolimits_{{t,:}}-FlashSoftmaxGrad \left( dO,O \left) \left) \right] \right. \right. \right. \right. $$阶段3计算$dQ$与$dK$:$$ d\mathop{{Q}}\nolimits_{{t,:}}d\mathop{{S}}\nolimits_{{t,:}}K \left[ u \left] \mathop{{}}\nolimits_{{:t,:}}/\sqrt{{d\mathop{{}}\nolimits_{{k,:}}}}\right. \right. $$$$ dK \left[ u \left] \mathop{{}}\nolimits_{{:t,:}}dS\mathop{{}}\nolimits_{{t,:t}}\mathop{{}}\nolimits^{{T}}\text{}Q/\sqrt{{d\mathop{{}}\nolimits_{{t,:}}}}\right. \right. $$参数说明参数名输入/输出/属性描述数据类型数据格式query输入attention结构的输入Q。BFLOAT16、FLOAT16NDkey输入attention结构的输入K。BFLOAT16、FLOAT16NDvalue输入attention结构的输入v。BFLOAT16、FLOAT16NDsparseIndices输入稀疏场景下选择的权重较高的注意力索引。INT32NDdOut输入注意力输出矩阵的梯度。BFLOAT16、FLOAT16NDout输入注意力输出矩阵。BFLOAT16、FLOAT16NDsoftmaxMax输入注意力正向计算的中间输出。FLOAT32NDsoftmaxSum输入注意力正向计算的中间输出。FLOAT32NDactualSeqLengthsQueryOptional输入每个Batch中Query的有效token数。INT32NDactualSeqLengthskvOptional输入每个Batch中Key、value的有效token数。INT32NDqueryRopeOptional输入MLA rope部分Query位置编码的输出。BFLOAT16、FLOAT16NDkeyRopeOptional输入MLA rope部分Key位置编码的输出。BFLOAT16、FLOAT16NDscaleValue属性缩放系数。FLOAT32-sparseBlockSize属性选择的块的大小。INT64-layout属性layout格式。STRING-sparseMode属性sparse的模式。INT64-preTokens属性Attention算子里, 对S矩阵的滑窗起始位置。INT64-nextTokens属性Attention算子里, 对S矩阵的滑窗终止位置。INT64-deterministic属性确定性计算。BOOL-dQuery输出表示query的梯度。BFLOAT16、FLOAT16NDdKey输出表示key的梯度。BFLOAT16、FLOAT16NDdValue输出表示value的梯度。BFLOAT16、FLOAT16NDdQueryRopeOptional输出表示queryRope的梯度。BFLOAT16、FLOAT16NDdKeyRopeOptional输出表示keyRope的梯度。BFLOAT16、FLOAT16ND约束说明参数query中的D和key、value的D值相等为512参数query_rope中的Dr和key_rope的Dr值相等为64。参数query、key、value的数据类型必须保持一致。当前只支持value和key完全一致的场景。当前仅支持sparseMode0或3无mask或以右顶点为划分的下三角场景仅支持BSND或TND layout关于数据shape的约束如下B取值范围1~256。S1、S21~128KS1、S2支持不等长。N1支持1/2/4/8/16/32/64/128。Ascend 950PR/Ascend 950DT 额外还支持48、24、12、6、3。N2仅支持1。D仅支持512。Drope仅支持64。topk1024、2048、3072、4096、5120、6144、7168、8192。不建议topk * sparseBlockSize超过100k由于内部算法硬件限制可能会导致oom。确定性计算SparseFlashAttentionGrad默认非确定性实现支持通过aclrtCtxSetSysParamOpt开启确定性。调用示例调用方式样例代码说明aclnn接口test_aclnn_sparse_flash_attention_grad通过 aclnnSparseFlashAttentionGrad 接口方式调用算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考