昇腾社区首页
中文
注册
开发者
下载

aclnnRingAttentionUpdateV2

产品支持情况

[object Object]undefined

功能说明

  • 算子功能:将两次FlashAttention的输出根据其不同的softmax的max和sum更新。接口的区别是:在输入排布为TND的场景下,原RingAttentionUpdate接口中的softmax相关输入的数据排布为BNS8,RingAttentionUpdateV2 接口支持传入字符串参数 inputSoftmaxLayout,用于控制 softmax 相关输入的数据排布是否与attention保持一致(即采用 TND 排布)。
  • 计算公式:
softmax_max=max(prev_softmax_max,cur_softmax_max)softmax\_max = max(prev\_softmax\_max, cur\_softmax\_max) softmax_sum=prev_softmax_sumexp(prev_softmax_maxsoftmax_max)+cur_softmax_sumexp(cur_softmax_maxsoftmax_max)softmax\_sum = prev\_softmax\_sum * exp(prev\_softmax\_max - softmax\_max) + cur\_softmax\_sum * exp(cur\_softmax\_max - softmax\_max) attn_out=prev_attn_outexp(prev_softmax_maxsoftmax_max)prev_softmax_sum/softmax_sum+cur_attn_outexp(cur_softmax_maxsoftmax_max)cur_softmax_sum/softmax_sumattn\_out = prev\_attn\_out * exp(prev\_softmax\_max - softmax\_max) * prev\_softmax\_sum / softmax\_sum + cur\_attn\_out * exp(cur\_softmax\_max - softmax\_max) * cur\_softmax\_sum / softmax\_sum

函数原型

  • 每个算子分为,必须先调用“aclnnRingAttentionUpdateV2GetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnRingAttentionUpdateV2”接口执行计算。
[object Object]
[object Object]

aclnnRingAttentionUpdateV2GetWorkspaceSize

  • 参数说明:

    [object Object]
  • 返回值:

    aclnnStatus:返回状态码,具体参见

    第一段接口完成入参校验,出现以下场景时报错:

    [object Object]

aclnnRingAttentionUpdateV2

  • 参数说明:

    [object Object]
  • 返回值:

    aclnnStatus:返回状态码,具体参见

约束说明

  • 确定性计算:
    • aclnnRingAttentionUpdateV2默认确定性实现。
  • 当inputLayoutOptional为“TND”时,prevAttnOut的最后一个维度需要为64的倍数。
  • 当inputLayoutOptional为“TND”时,actualSeqQlenOptional为必填。
  • 当inputLayoutOptional为“TND”时,请注意N和D的大小,限制为:N<=256, D<=768。
  • 当inputLayoutOptional为“TND”时,inputSoftmaxLayoutOptional才生效。inputSoftmaxLayoutOptional只支持三种输入:空字符串、“SBH”、“TND”

调用示例

示例代码如下,仅供参考,具体编译和执行过程请参考

[object Object]