CANN稀疏FlashAttention反向算子

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、FLOAT16ND
key输入attention结构的输入K。BFLOAT16、FLOAT16ND
value输入attention结构的输入v。BFLOAT16、FLOAT16ND
sparseIndices输入稀疏场景下选择的权重较高的注意力索引。INT32ND
dOut输入注意力输出矩阵的梯度。BFLOAT16、FLOAT16ND
out输入注意力输出矩阵。BFLOAT16、FLOAT16ND
softmaxMax输入注意力正向计算的中间输出。FLOAT32ND
softmaxSum输入注意力正向计算的中间输出。FLOAT32ND
actualSeqLengthsQueryOptional输入每个Batch中，Query的有效token数。INT32ND
actualSeqLengthskvOptional输入每个Batch中，Key、value的有效token数。INT32ND
queryRopeOptional输入MLA rope部分：Query位置编码的输出。BFLOAT16、FLOAT16ND
keyRopeOptional输入MLA rope部分：Key位置编码的输出。BFLOAT16、FLOAT16ND
scaleValue属性缩放系数。FLOAT32-
sparseBlockSize属性选择的块的大小。INT64-
layout属性layout格式。STRING-
sparseMode属性sparse的模式。INT64-
preTokens属性Attention算子里, 对S矩阵的滑窗起始位置。INT64-
nextTokens属性Attention算子里, 对S矩阵的滑窗终止位置。INT64-
deterministic属性确定性计算。BOOL-
dQuery输出表示query的梯度。BFLOAT16、FLOAT16ND
dKey输出表示key的梯度。BFLOAT16、FLOAT16ND
dValue输出表示value的梯度。BFLOAT16、FLOAT16ND
dQueryRopeOptional输出表示queryRope的梯度。BFLOAT16、FLOAT16ND
dKeyRopeOptional输出表示keyRope的梯度。BFLOAT16、FLOAT16ND

约束说明

  • 参数query中的D和key、value的D值相等为512，参数query_rope中的Dr和key_rope的Dr值相等为64。
  • 参数query、key、value的数据类型必须保持一致。
  • 当前只支持value和key完全一致的场景。
  • 当前仅支持sparseMode=0或3（无mask或以右顶点为划分的下三角场景）
  • 仅支持BSND或TND layout；关于数据shape的约束如下：
    • B：取值范围1~256。
    • S1、S2：1~128K；S1、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。
    • topk：1024、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），仅供参考