[object Object][object Object][object Object]undefined[object Object]

实现“Transformer Attention Score”的融合计算，实现的计算公式如下：

[object Object][object Object]

• query（[object Object]）：数据类型支持[object Object]、[object Object]、[object Object]，数据格式支持$ND$。综合约束请见。

• key（[object Object]）：数据类型支持[object Object]、[object Object]、[object Object]，数据格式支持$ND$。综合约束请见。

• value（[object Object]）：数据类型支持[object Object]、[object Object]、[object Object]，数据格式支持$ND$。综合约束请见。

• head_num（[object Object]）：代表head个数，数据类型支持[object Object]。综合约束请见。

• input_layout（[object Object]）：代表输入[object Object]、[object Object]、[object Object]的数据排布格式，支持BSH、SBH、BSND、BNSD、TND（[object Object]/[object Object]需传值，[object Object]t为TND时即为varlen场景）；后续章节如无特殊说明，S表示[object Object]或[object Object]、[object Object]的sequence length，Sq表示query的sequence length，Skv表示[object Object]、[object Object]的sequence length，SS表示Sq*Skv。

• pse（[object Object]）：可选参数，表示位置编码。数据类型支持[object Object]、[object Object]、[object Object]，数据格式支持$ND$。非varlen场景支持四维输入，包含BNSS格式、BN1Skv格式、1NSS格式。如果非varlen场景Sq大于1024或varlen场景、每个batch的Sq与Skv等长且是sparse_mode为0、2、3的下三角掩码场景，可使能alibi位置编码压缩，此时只需要输入原始PSE最后1024行进行内存优化，即alibi_compress = ori_pse[:, :, -1024:, :]，参数每个batch不相同时，输入BNHSkv(H=1024)，每个batch相同时，输入1NHSkv(H=1024)。

• padding_mask（[object Object]）：暂不支持该传参。

• atten_mask（[object Object]）：可选参数，取值为1代表该位不参与计算（不生效），为0代表该位参与计算，数据类型支持[object Object]、[object Object]，数据格式支持$ND$，输入shape类型支持BNSS格式、B1SS格式、11SS格式、SS格式。varlen场景只支持SS格式，SS分别是maxSq和maxSkv。综合约束请见。

• scale（[object Object]）：可选参数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持[object Object]，默认值为1。

• keep_prob（[object Object]）：可选参数，代表Dropout中1的比例，取值范围为(0, 1] 。数据类型支持[object Object]，默认值为1，表示全部保留。

• pre_tockens（[object Object]）：用于稀疏计算的参数，可选参数，数据类型支持[object Object]，默认值为2147483647。综合约束请见。

• next_tockens（[object Object]）：用于稀疏计算的参数，可选参数，数据类型支持[object Object]，默认值为2147483647。[object Object]和[object Object]取值与[object Object]的关系请参见[object Object]参数，参数取值与[object Object]分布不一致会导致精度问题。综合约束请见。

• inner_precise（[object Object]）：用于提升精度，数据类型支持[object Object]，默认值为0。

  [object Object]

• prefix（[object Object]）：可选参数，代表prefix稀疏计算场景每个Batch的N值。数据类型支持[object Object]，数据格式支持$ND$。综合约束请见。

• actual_seq_qlen（[object Object]）：可选参数，varlen场景时需要传入此参数。表示[object Object]每个S的累加和长度，数据类型支持[object Object]，数据格式支持$ND$。综合约束请见。

  比如真正的S长度列表为：2 2 2 2 2，则[object Object]传：2 4 6 8 10。

• actual_seq_kvlen（[object Object]）：可选参数，varlen场景时需要传入此参数。表示[object Object]/[object Object]每个S的累加和长度。数据类型支持[object Object]，数据格式支持$ND$。综合约束请见。

  比如真正的S长度列表为：2 2 2 2 2，则actual_seq_kvlen传：2 4 6 8 10。

• sparse_mode（[object Object]）：表示sparse的模式，可选参数，默认值为0。取值如所示，不同模式的原理参见。当整网的[object Object]都相同且shape小于2048*2048时，建议使用defaultMask模式，来减少内存使用量。综合约束请见。

  表1 sparse_mode不同取值场景说明

  [object Object][object Object]

  [object Object]

• gen_mask_parallel（[object Object]）：DSA生成dropout随机数向量mask的控制开关。默认值为True：同AI Core并行计算；设为False：同AI Core串行计算。

• sync（[object Object]）：DSA生成dropout随机数向量mask的控制开关。默认值为False：dropout mask异步生成；设为True：dropout mask同步生成。

• softmax_layout（[object Object]）：可选参数，用于控制TND场景下softmax的输出（softmax_max和softmax_sum）的数据排布方式。当前仅在input_layout=“TND”时进行配置，仅支持传入“TND”。默认情况下，softmax的输出排布为NTD排布；传入TND时，softmax的输出排布为TND排布。

[object Object]

共7个输出，类型依次为Tensor、Tensor、Tensor、Tensor、int、int、int。

• 第1个输出为[object Object]，计算公式的最终输出$attention\_out$，数据类型支持[object Object]、[object Object]、[object Object]。
• 第2个输出为[object Object]，Softmax计算的Max中间结果，用于反向计算，数据类型支持[object Object]。
• 第3个输出为[object Object]，Softmax计算的Sum中间结果，用于反向计算，数据类型支持[object Object]。
• 第4个输出为[object Object]，预留参数，暂未使用。
• 第5个输出为[object Object]，DSA生成dropout mask中，Philox算法的seed。
• 第6个输出为[object Object]，DSA生成dropout mask中，Philox算法的offset。
• 第7个输出为[object Object]，DSA生成dropout mask的长度。

[object Object]

• 该接口仅在训练场景下使用。

• 该接口暂不支持图模式。

• 输入[object Object]、[object Object]、[object Object]、[object Object]的数据类型必须一致。

• 输入[object Object]、[object Object]、[object Object]的[object Object]必须一致。

• 输入[object Object]、[object Object]、[object Object]的shape说明：

  • 输入[object Object]和[object Object]的shape必须一致。
  • B：batchsize必须相等；非varlen场景B取值范围1~2M；varlen场景B取值范围1~2K。
  • D：Head Dim必须满足Dq=Dk和Dk≥Dv，取值范围1~768。
  • S：sequence length，取值范围1~1M。

• varlen场景下：

  • 要求T（B*S）取值范围1~1M。
  • [object Object]输入不支持补pad，即[object Object]中不能存在某一行全1的场景。

• 支持输入[object Object]的N和[object Object]/[object Object]的N不相等，但必须成比例关系，即Nq/Nkv必须是非0整数，Nq取值范围1~256。当Nq/Nkv > 1时，即为GQA(grouped-query attention)；当Nq/Nkv=1时，即为MHA(multi-head attention)。

  [object Object]

• [object Object]取值说明：

  • [object Object]为1、2、3、4、5、6、7、8时，应传入对应正确的[object Object]，否则将导致计算结果错误。当[object Object]输入为None时，[object Object]，[object Object]，[object Object]参数不生效，固定为全计算。
  • [object Object]配置为1、2、3、5、6时，用户配置的[object Object]、[object Object]不会生效。
  • [object Object]配置为0、4时，需保证[object Object]与[object Object]、[object Object]的范围一致。
  • [object Object]配置为7或者8时，不支持可选参数[object Object]。

• [object Object]稀疏计算场景B不大于32，varlen场景不支持非压缩prefix，即不支持sparse_mode=5；当Sq>Skv时，[object Object]的N值取值范围[0, Skv]，当Sq<=Skv时，[object Object]的N值取值范围[Skv-Sq, Skv]。

• 支持[object Object]中某个Batch上的S长度为0；如果存在S为0的情况，不支持[object Object]输入，假设真实的S长度为[2, 2, 0, 2, 2]，则传入的[object Object]为[2, 4, 4, 6, 8]。[object Object]的长度取值范围为1~2K，varlen场景下长度最大支持1K。

• TND格式下，支持尾部部分Batch不参与计算，此时[object Object]和[object Object]尾部传入对应个数个0即可。假设真实的S长度为[2, 3, 4, 5, 6]，此时后两个Batch不参与计算，则传入的[object Object]为[2, 5, 9, 0, 0]。

• 部分场景下，如果计算量过大可能会导致算子执行超时(aicore error类型报错，errorStr为：timeout or trap error)，此时建议做轴切分处理，注：这里的计算量会受B、S、N、D等参数的影响，值越大计算量越大。

[object Object]

单算子模式调用：

[object Object]

[object Object]的工作原理为，在Mask为True的位置遮蔽[object Object](Q)与[object Object](K)的转置矩阵乘积的值，示意如下：

QK[object Object]T[object Object]矩阵在[object Object]为True的位置会被遮蔽，效果如下：

说明：下图中的蓝色表示保留该值，[object Object]中，应该配置为False；阴影表示遮蔽该值，[object Object]中应配置为True。

• 当[object Object]为0时，代表defaultMask模式。

  • 不传mask：如果[object Object]未传入则不做mask操作，[object Object]取值为None，忽略[object Object]和[object Object]取值。Masked QK[object Object]T[object Object]矩阵示意如下：

    

  • [object Object]取值为0，[object Object]大于等于Sq，表示causal场景sparse，[object Object]应传入下三角矩阵，此时[object Object]和[object Object]之间的部分需要计算，Masked QK[object Object]T[object Object]矩阵示意如下：

    

    [object Object]应传入下三角矩阵，示意如下：

    

  • [object Object]小于Sq，[object Object]小于Skv，且都大于等于0，表示band场景，此时[object Object]和[object Object]之间的部分需要计算。Masked QK[object Object]T[object Object]矩阵示意如下：

    

    [object Object]应传入band形状矩阵，示意如下：

    

  • [object Object]为负数，以pre_tockens=9，next_tockens=-3为例，[object Object]和[object Object]之间的部分需要计算。Masked QK[object Object]T[object Object]示意如下：

    说明：[object Object]为负数时，[object Object]取值必须大于等于[object Object]的绝对值，且[object Object]的绝对值小于Skv。

    

  • [object Object]为负数，以next_tockens=7，pre_tockens=-3为例，[object Object]和[object Object]之间的部分需要计算。Masked QK[object Object]T[object Object]示意如下：

    说明：[object Object]为负数时，[object Object]取值必须大于等于[object Object]的绝对值，且[object Object]的绝对值小于Sq。

    

• 当[object Object]为1时，代表allMask，即传入完整的[object Object]矩阵。

  该场景下忽略[object Object]、[object Object]取值，Masked QK[object Object]T[object Object]矩阵示意如下：

  

• 当[object Object]为2时，代表leftUpCausal模式的mask，对应以左上顶点划分的下三角场景（参数起点为左上角）。该场景下忽略[object Object]、[object Object]取值，Masked QK[object Object]T[object Object]矩阵示意如下：

  

  传入的atten_mask为优化后的压缩下三角矩阵（2048*2048），压缩下三角矩阵示意（下同）：

  

• 当[object Object]为3时，代表rightDownCausal模式的mask，对应以右下顶点划分的下三角场景（参数起点为右下角）。该场景下忽略[object Object]、[object Object]取值。[object Object]为优化后的压缩下三角矩阵（2048*2048），Masked QK[object Object]T[object Object]矩阵示意如下：

  

• 当[object Object]为4时，代表band场景，即计算[object Object]和[object Object]之间的部分，参数起点为右下角，[object Object]和[object Object]之间需要有交集。[object Object]为优化后的压缩下三角矩阵（2048*2048）。Masked QK[object Object]T[object Object]矩阵示意如下：

  

• 当[object Object]为5时，代表prefix非压缩场景，即在rightDownCasual的基础上，左侧加上一个长为Sq，宽为N的矩阵，N的值由可选参数prefix获取，例如下图中表示batch=2场景下prefix传入数组[4,5]，每个batch轴的N值可以不一样，参数起点为左上角。

  该场景下忽略[object Object]、[object Object]取值，[object Object]矩阵数据格式须为BNSS或B1SS，Masked QK[object Object]T[object Object]矩阵示意如下：

  

  [object Object]应传入矩阵示意如下：

  

• 当[object Object]为6时，代表prefix压缩场景，即prefix场景时，attenMask为优化后的压缩下三角+矩形的矩阵（3072*2048）：其中上半部分[2048，2048]的下三角矩阵，下半部分为[1024,2048]的矩形矩阵，矩形矩阵左半部分全0，右半部分全1，[object Object]应传入矩阵示意如下。该场景下忽略[object Object]、[object Object]取值。

  

• 当[object Object]为7时，表示varlen且为长序列外切场景（即长序列在模型脚本中进行多卡切query的sequence length）；用户需要确保外切前为使用sparse_mode=3的场景；当前mode下用户需要设置[object Object]和[object Object]（起点为右下顶点），且需要保证参数正确，否则会存在精度问题。

  Masked QK[object Object]T[object Object]矩阵示意如下，在第二个batch对[object Object]进行切分，[object Object]和[object Object]不切分，4x6的mask矩阵被切分成2x6和2x6的mask，分别在卡1和卡2上计算：

  • 卡1的最后一块mask为band类型的mask，配置pre_tockens=6（保证大于等于最后一个Skv），next_tockens=-2，[object Object]应传入{3,5}，[object Object]应传入{3,9}。
  • 卡2的mask类型切分后不变，[object Object]为3，[object Object]应传入{2,7,11}，[object Object]应传入{6,11,15}。

  

  [object Object]

• 当[object Object]为8时，表示varlen且为长序列外切场景；用户需要确保外切前为使用sparse_mode=2的场景；当前mode下用户需要设置[object Object]和[object Object]（起点为右下顶点），且需要保证参数正确，否则会存在精度问题。

  Masked QK[object Object]T[object Object]矩阵示意如下，在第二个batch对[object Object]进行切分，[object Object]和[object Object]不切分，5x4的mask矩阵被切分成2x4和3x4的mask，分别在卡1和卡2上计算：

  • 卡1的mask类型切分后不变，[object Object]为2，[object Object]应传入{3,5}，[object Object]应传入{3,7}。
  • 卡2的第一块mask为band类型的mask，配置pre_tockens=4（保证大于等于第一个Skv），next_tockens=1，[object Object]应传入{3,8,12}，[object Object]应传入{4,9,13}。

  

  [object Object]