接口功能:训练场景下,使用FlashAttention算法实现self-attention(自注意力)的计算。对标竞品适配gptoss模型支持sink功能。该接口query、key、value参数支持多个长度相等或者多个长度不相等的sequence
- 该接口相较于接口,功能差异如下::
- 调整Dropout功能:在keepProb小于1.0时,若没有外部传入的DropoutMask,则使用新增参数seed和offset生成DropoutMask;若有外部传入的DropoutMask,则使用外部传入的DropoutMask。 -增加了sinkOptional可选输入,数据类型支持FLOAT32。
- 该接口相较于接口,功能差异如下:
- 调整Dropout功能:在keepProb小于1.0时,若没有外部传入的DropoutMask,则使用新增参数seed和offset生成DropoutMask;若有外部传入的DropoutMask,则使用外部传入的DropoutMask。
- 该接口相较于接口,功能差异如下::
计算公式:
注意力的正向计算公式如下:
其中增加sink之后计算逻辑见下,主要修改相关softmax_max和softmax_sum逻辑计算部分
每个算子分为,必须先调用“aclnnFlashAttentionScoreV4GetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnFlashAttentionScoreV4”接口执行计算。
确定性计算:
- aclnnFlashAttentionScoreV4默认确定性实现。
输入query、key、value的
- B:batchsize必须相等。
- D:Head-Dim必须满足(qD == kD && kD >= vD)。
- inputLayout必须一致。
输入queryRopeOptional与query的输入shape仅在D维度不同,其他shape参数应该相同
输入keyRopeOptional与key的输入shape仅在D维度不同,其他shape参数应该相同
关于数据shape的约束,以inputLayout的TND、BSND、BNSD为例(BSH、SBH下H=N*D),其中:
- T(B*S):取值范围为1~1M;TND格式下actualSeqQLenOptional支持的最大长度为20000。
- B:取值范围为1~2M。带prefixOptional的时候B最大支持2K。
- N:取值范围为1~256。
- S:取值范围为1~1M。
- D:取值范围为1~768。
query、key、value数据排布格式支持从多种维度解读,其中B(Batch)表示输入样本批量大小、S(Seq-Length)表示输入样本序列长度、H(Head-Size)表示隐藏层的大小、N(Head-Num)表示多头数、D(Head-Dim)表示隐藏层最小的单元尺寸,且满足D=H/N。
innerPrecise: 当前0、1为保留配置值,2为开启无效行计算,其功能是避免在计算过程中存在整行mask进而导致精度有损失,但是该配置会导致性能下降。 如果算子可判断出存在无效行场景,会自动开启无效行计算,例如sparseMode为3,Sq > Skv场景。
pseType 各个取值含义
[object Object]undefined
pseType为2或3的时候,当前只支持Sq和Skv等长。
sparseMode约束如下:
部分场景下,如果计算量过大可能会导致算子执行超时(aicore error类型报错,errorStr为:timeout or trap error),此时建议做轴切分处理,注:这里的计算量会受B、S、N、D等参数的影响,值越大计算量越大。
band场景,preTokens和nextTokens之间必须要有交集。
prefixOptional稀疏计算场景,场景包括sequence长度相等的场景下sparseMode=5、sparseMode=6;sequence长度不相等的场景下sparseMode=6。这两种场景下,当Sq > Skv时,prefix的N值取值范围[0, Skv];当Sq <= Skv时,prefix的N值取值范围[Skv-Sq, Skv]。当sparseModeOptional=5、prefix的N > Skv或prefixOptional不传时执行全计算,sparseModeOptional=6要求prefixOptional必传。
realShiftOptional:Sq大于1024时如果配置BNHS、1NHS,需要Sq和Skv等长。
actualSeqQLenOptional输入支持某个Batch上的S长度为0,此时不支持可选输入realShiftOptional。
attenMaskOptional输入不支持补pad,即attenMaskOptional中不能存在某一行全1的场景。
支持actualSeqQLenOptional中某个Batch上的S长度为0;如果存在S为0的情况,不支持pse输入, 假设真实的S长度为[2,2,0,2,2],则传入的actualSeqQLenOptional为[2,4,4,6,8]。
Atlas 350 加速卡:
- seed和offset只在keepProb小于1.0时生效,否则不生效。
- keepProb小于1.0时,若dropMaskOptional非nullptr,则使用输入的dropMask;否则使用seed和offset生成的dropMask。
TND格式下,支持尾部部分Batch不参与计算,此时actual_seq_q_len和actual_seq_kv_len尾部传入对应个数的0即可。假设真实S长度为[2, 3, 4, 5, 6],若希望最后两个Batch不参与计算,则传入的actual_seq_q_len为[2, 3, 4, 0, 0]。此时若需要传入prefixOptional,其尾部也需要传入同等数量的0,例如[1, 1, 1, 0, 0]。
softmaxOutLayout支持传入:空字符串、"same_as_input"。