Skip to main content

aclnnFlashAttentionScore

aclnnPromptFlashAttention​

接口原型​

每个算子有两段接口，必须先调用“aclnnPromptFlashAttentionGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnPromptFlashAttention”接口执行计算。两段式接口如下：

  • **第一段接口：**aclnnStatus aclnnPromptFlashAttentionGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *paddingMask, const aclTensor *attenMask, const aclIntArray *actualSeqLengths, int64_t numHeads, double scaleValue, int64_t preTokens, int64_t nextTokens, char* inputLayout, int64_t numKeyValueHeads, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)
  • **第二段接口：****aclnnStatus aclnnPromptFlashAttention(**void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述​

  • 算子功能：全量推理场景的FlashAttention算子。

  • 计算公式：

    self-attention（自注意力）利用输入样本自身的关系构建了一种注意力模型。其原理是假设有一个长度为n的输入样本序列x，x的每个元素都是一个d维向量，可以将每个d维向量看作一个token embedding，将这样一条序列经过3个权重矩阵变换得到3个维度为n*d的矩阵。

    self-attention的计算公式一般定义如下，其中Q、K、V为输入样本的重要属性元素，是输入样本经过空间变换得到，且可以统一到一个特征空间中。

    本算子中Score函数采用Softmax函数，self-attention计算公式为

    其中Q和KT的乘积代表输入x的注意力，为避免该值变得过大，通常除以d的开根号进行缩放，并对每行进行softmax归一化，与V相乘后得到一个n*d的矩阵。

aclnnPromptFlashAttentionGetWorkspaceSize​

  • 接口定义：

    aclnnStatus aclnnPromptFlashAttentionGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *paddingMask, const aclTensor *attenMask, const aclIntArray *actualSeqLengths, int64_t numHeads, double scaleValue, int64_t preTokens, int64_t nextTokens, char* inputLayout, int64_t numKeyValueHeads, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：

    • query：Device侧的aclTensor，公式中的输入Q，数据类型支持FLOAT16，不支持非连续的Tensor，数据格式支持ND。

    • key：Device侧的aclTensor，公式中的输入K，数据类型支持FLOAT16，不支持非连续的Tensor，数据格式支持ND。

    • value：Device侧的aclTensor，公式中的输入V，数据类型支持FLOAT16，不支持非连续的Tensor，数据格式支持ND。

    • paddingMask：Device侧的aclTensor，可选参数，数据类型支持FLOAT16，不支持非连续的Tensor，数据格式支持ND。

    • attenMask：Device侧的aclTensor，可选参数，代表下三角全为0上三角全为负无穷的倒三角mask矩阵，不支持非连续的Tensor，数据类型支持BOOL、FLOAT16，数据格式支持ND。

    • actualSeqLengths：Host侧的aclIntArray，可选参数，代表不同Batch中query的Sequence Length，数据类型支持INT64。

    • numHeads：Host侧的int，代表head个数，数据类型支持INT64。

    • scaleValue：Host侧的double，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。

    • preTokens：Host侧的int，可选属性，用于稀疏计算，数据类型支持INT64。

    • nextTokens：Host侧的int，可选属性，用于稀疏计算，数据类型支持INT64。

    • inputLayout：Host侧的字符指针，用于标识输入query、key、value的数据排布格式，当前支持BSH、BSND。

      :::note 说明 query、key、value数据排布格式支持从多种维度解读，其中B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。 :::

    • numKeyValueHeads：Host侧的int，可选属性，代表key、value中head个数，用于支持GQA（Grouped-Query Attention，分组查询注意力）场景，默认为0，表示和query的head个数相等，数据类型支持INT64。

    • attentionOut：Device侧的aclTensor，公式中的输出，数据类型支持FLOAT16，数据格式支持ND。

    • workspaceSize：返回用户需要在Device侧申请的workspace大小。

    • executor：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    :::note 说明 第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的query、key、value、attentionOut是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：query、key、value、paddingMask、attenMask、attentionOut的数据类型和数据格式不在支持的范围内。 :::

aclnnPromptFlashAttention​

  • 接口定义：

    aclnnStatus aclnnPromptFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：

    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnPromptFlashAttentionGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束与限制​

无

父主题： 融合类算子接口

aclnnIncreFlashAttention​

接口原型​

每个算子有两段接口，必须先调用“aclnnIncreFlashAttentionGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnIncreFlashAttention”接口执行计算。两段式接口如下：

  • **第一段接口：**aclnnStatus aclnnIncreFlashAttentionGetWorkspaceSize(const aclTensor *query, const aclTensorList *key, const aclTensorList *value, const aclTensor *paddingMask, const aclTensor *attenMask, const aclIntArray *actualSeqLengths, int64_t numHeads, double scaleValue, char *inputLayout, int64_t numKeyValueHeads, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)
  • 第二段接口：aclnnstatus aclnnIncreFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述​

  • 算子功能：对于自回归（Auto-regressive）的语言模型，随着新词的生成，推理输入长度不断增大。在原来全量推理的基础上实现****增量推理，query的S轴固定为1，key和value是经过KV Cache后，将之前推理过的state信息，叠加在一起，每个Batch对应S轴的实际长度可能不一样，输入的数据是经过padding后的固定长度数据。

    相比全量场景的FlashAttention算子（aclnnPromptFlashAttention），增量推理的流程与正常全量推理并不完全等价，不过增量推理的精度并无明显劣化。

    :::note 说明 KV Cache是大模型推理性能优化的一个常用技术。采样时，Transformer模型会以给定的prompt/context作为初始输入进行推理（可以并行处理），随后逐一生成额外的token来继续完善生成的序列（体现了模型的自回归性质）。在采样过程中，Transformer会执行自注意力操作，为此需要给当前序列中的每个项目（无论是prompt/context还是生成的token）提取键值（KV）向量。这些向量存储在一个矩阵中，通常被称为kv缓存（KV Cache）。 :::

  • 计算公式：

    self-attention（自注意力）利用输入样本自身的关系构建了一种注意力模型。其原理是假设有一个长度为n的输入样本序列x，x的每个元素都是一个d维向量，可以将每个d维向量看作一个token embedding，将这样一条序列经过3个权重矩阵变换得到3个维度为n*d的矩阵。

    self-attention的计算公式一般定义如下，其中Q、K、V为输入样本的重要属性元素，是输入样本经过空间变换得到，且可以统一到一个特征空间中。

    本算子中Score函数采用Softmax函数，self-attention计算公式为

    其中Q和KT的乘积代表输入x的注意力，为避免该值变得过大，通常除以d的开根号进行缩放，并对每行进行softmax归一化，与V相乘后得到一个n*d的矩阵。

aclnnIncreFlashAttentionGetWorkspaceSize​

  • 接口定义：

    aclnnStatus aclnnIncreFlashAttentionGetWorkspaceSize(const aclTensor *query, const aclTensorList *key, const aclTensorList *value, const aclTensor *paddingMask, const aclTensor *attenMask, const aclIntArray *actualSeqLengths, int64_t numHeads, double scaleValue, char *inputLayout, int64_t numKeyValueHeads, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：

    • query：Device侧的aclTensor，公式中的输入Q，数据类型支持FLOAT16，数据格式支持ND。

    • key：Device侧的aclTensorList，公式中的输入K，数据类型支持FLOAT16，数据格式支持ND。

    • value：Device侧的aclTensorList，公式中的输入V，数据类型支持FLOAT16，数据格式支持ND。

    • paddingMask：Device侧的aclTensor，可选参数，数据类型支持FLOAT16，数据格式支持ND。

    • attenMask：Device侧的aclTensor，可选参数，数据类型支持BOOL、FLOAT16，数据格式支持ND。

    • actualSeqLengths：Host侧的aclIntArray，可选参数，数据类型支持INT64。

    • numHeads：Host侧的int64_t，代表head个数，数据类型支持INT64。

    • scaleValue：Host侧的double，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。

    • inputLayout：Host侧的字符指针，用于标识输入query、key、value的数据排布格式，当前支持BSH、BNSD。

      :::note 说明 query、key、value数据排布格式支持从多种维度解读，其中B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。 :::

    • numKeyValueHeads：Host侧的int64_t，可选参数，代表key、value中head个数，用于支持GQA（Grouped-Query Attention，分组查询注意力）场景，默认为0，表示和query的head个数相等，数据类型支持INT64。

    • attentionOut：Device侧的aclTensor，公式中的输出，数据类型支持FLOAT16，数据格式支持ND。

    • workspaceSize：返回用户需要在Device侧申请的workspace大小。

    • executor：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    :::note 说明 第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的query、key、value、attentionOut是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：query、key、value、paddingMask、attenMask、attentionOut的数据类型和数据格式不在支持的范围内。 :::

aclnnIncreFlashAttention​

  • 接口定义：

    aclnnstatus aclnnIncreFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：

    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnIncreFlashAttentionGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束与限制​

  • 参数key、value 中对应tensor的shape需要完全一致；非连续场景下 key、value 的tensorlist中的tensor的shape也需要完全一致，且batch只能为1。
  • 参数query和attentionOut的shape需要完全一致。
  • 参数query中的N和numHeads值相等，key、value的N和numKeyValueHeads值相等，并且numHeads是numKeyValueHeads的倍数关系。
  • 非连续场景下，参数key、value的tensorlist中tensor的个数等于query的B，shape需要完全一致。

父主题： 融合类算子接口

aclnnFlashAttentionScore​

接口原型​

每个算子有两段接口，必须先调用“aclnnFlashAttentionScoreGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnFlashAttentionScore”接口执行计算。两段式接口如下：

  • **第一段接口：**aclnnStatus aclnnFlashAttentionScoreGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShift, const aclTensor *dropMask, const aclTensor *paddingMask, const aclTensor *attenMask, double scaleValue, double keepProb, int64_t preTockens, int64_t nextTockens, int64_t headNum, string *inputLayout, int32_t innerPrecise, const aclTensor *softmaxMax, const aclTensor *softmaxSum, const aclTensor *softmaxOut, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)
  • **第二段接口：**aclnnStatus aclnnFlashAttentionScore(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述​

  • 算子功能：训练场景下，使用FlashAttention算法实现self-attention（自注意力）的计算。

  • 计算公式：

    注意力的正向计算公式如下：

aclnnFlashAttentionScoreGetWorkspaceSize​

  • 接口定义：

    aclnnStatus aclnnFlashAttentionScoreGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *realShift, const aclTensor *dropMask, const aclTensor *paddingMask, const aclTensor *attenMask, double scaleValue, double keepProb, int64_t preTockens, int64_t nextTockens, int64_t headNum, string *inputLayout, int32_t innerPrecise, const aclTensor *softmaxMax, const aclTensor *softmaxSum, const aclTensor *softmaxOut, const aclTensor *attentionOut, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：

    • query：Device侧的aclTensor，公式中输入Q，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • key：Device侧的aclTensor，公式中输入K，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • value：Device侧的aclTensor，公式中输入V，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • realShift：Device侧的aclTensor，公式中输入pse，可选参数，表示位置编码。数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • dropMask：Device侧的aclTensor，可选属性，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • paddingMask：Device侧的aclTensor，可选属性，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • attenMask：Device侧的aclTensor，可选属性，代表下三角全为0上三角全为负无穷的倒三角mask矩阵，数据类型支持BOOL，数据格式支持ND。

    • scaleValue：Host侧的double，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。

    • keepProb：Host侧的double，可选参数，代表dropMask中1的比例，数据类型支持DOUBLE。

    • preTockens：Host侧的int64_t，用于稀疏计算的参数，可选参数，数据类型支持INT64。

    • nextTockens：Host侧的int64_t，用于稀疏计算的参数，可选参数，数据类型支持INT64。

    • headNum：Host侧的int64_t，代表head个数，数据类型支持INT64。

    • inputLayout：Host侧的string，代表输入query、key、value的数据排布格式，支持BSH、SBH、BSND、BNSD。

      :::note 说明 query、key、value数据排布格式支持从多种维度解读，其中B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。 :::

    • innerPrecise：Host侧的int32_t，数据类型支持INT32，内部计算精度模式，其中0表示为高精度，1表示为高性能。

    • softmaxMax：Device侧的aclTensor，正向计算的中间输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • softmaxSum：Device侧的aclTensor，正向计算的中间输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • softmaxOut：Device侧的aclTensor，正向计算的softmax输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • attentionOut：Device侧的aclTensor，算子最终输出，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • workspaceSize：返回用户需要在Device侧申请的workspace大小。

    • executor：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    :::note 说明 第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的query、key、value、attention_out是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：query、key、value、realShift、dropMask、paddingMask、attenMask、softmaxMax、softmaxSum、softmaxOut、attentionOut的数据类型和数据格式不在支持的范围内。 :::

aclnnFlashAttentionScore​

  • 接口定义：

    aclnnStatus aclnnFlashAttentionScore(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：

    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnFlashAttentionScoreGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束与限制​

关于数据shape的约束，以inputLayout的BSND、BNSD为例（BSH、SBH下H=N*D），其中

  • B：取值范围为1~256。
  • N：取值范围为1~256。
  • S：取值范围为1~32K，且为16的倍数。
  • D：取值为64、96、128、256。

父主题： 融合类算子接口

aclnnFlashAttentionScoreGrad​

接口原型​

每个算子有两段接口，必须先调用“aclnnFlashAttentionScoreGradGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnFlashAttentionScoreGrad”接口执行计算。两段式接口如下：

  • **第一段接口：**aclnnStatus aclnnFlashAttentionScoreGradGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *dy, const aclTensor *pseShift, const aclTensor *dropMask, const aclTensor *paddingMask, const aclTensor *attenMask, const aclTensor *softmaxMax, const aclTensor *softmaxSum, const aclTensor *softmaxIn, const aclTensor *attentionIn, const aclTensor *dq, const aclTensor *dk, const aclTensor *dv, const aclTensor *dpse, double scaleValue, double keepProb, int64_t precTockens, int64_t nextTockens, int64_t headNum, string *inputLayout, int32_t innerPrecise, uint64_t *workspaceSize, aclOpExecutor **executor)
  • **第二段接口：**aclnnStatus aclnnFlashAttentionScoreGrad(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述​

  • 算子功能：训练场景下计算注意力的反向输出，即aclnnFlashAttentionScore的反向计算。

  • 计算公式：

    已知注意力的正向计算公式为

    为方便表达，以变量S和P表示计算公式：

    那么注意力的反向计算公式为

aclnnFlashAttentionGradGetWorkspaceSize​

  • 接口定义：

    aclnnStatus aclnnFlashAttentionScoreGradGetWorkspaceSize(const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *dy, const aclTensor *pseShift, const aclTensor *dropMask, const aclTensor *paddingMask, const aclTensor *attenMask, const aclTensor *softmaxMax, const aclTensor *softmaxSum, const aclTensor *softmaxIn, const aclTensor *attentionIn, const aclTensor *dq, const aclTensor *dk, const aclTensor *dv, const aclTensor *dpse，double scaleValue, double keepProb, int64_t preTockens, int64_t nextTockens, int64_t headNum, string *inputLayout, int32_t innerPrecise, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：

    • query：Device侧的aclTensor，公式中的输入Q，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • key：Device侧的aclTensor，公式中的输入K，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • value：Device侧的aclTensor，公式中的输入V，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • dy：Device侧的aclTensor，公式中的输入dY，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • pseShift：Device侧的aclTensor，公式中的输入pse，可选参数，表示位置编码。数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • dropMask：Device侧的aclTensor，可选属性，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • paddingMask：Device侧的aclTensor，可选属性，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • attenMask：Device侧的aclTensor，可选属性，代表下三角全为0上三角全为负无穷的倒三角mask矩阵，数据类型支持BOOL，数据格式支持ND。

    • softmaxMax：Device侧的aclTensor，注意力正向计算的中间输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • softmaxSum：Device侧的aclTensor，注意力正向计算的中间输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • softmaxIn：Device侧的aclTensor，注意力正向计算的中间输出，数据类型支持FLOAT、FLOAT16，数据格式支持ND。

    • attentionIn：Device侧的aclTensor，注意力正向计算的最终输出attentionOut，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • dq：Device侧的aclTensor，公式中的dQ，表示query的梯度，计算输出，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • dk：Device侧的aclTensor，公式中的dK，表示key的梯度，计算输出，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • dv：Device侧的aclTensor，公式中的dV，表示value的梯度，计算输出，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • dpse：Device侧的aclTensor，公式中的d(pse)，表示pse的梯度，计算输出，数据类型支持FLOAT、FLOAT16、BFLOAT16，数据格式支持ND。

    • scaleValue：Host侧的double，公式中d开根号的倒数，代表缩放系数，作为计算流中Muls的scalar值，数据类型支持DOUBLE。

    • keepProb：Host侧的double，可选参数，代表dropMask中1的比例，数据类型支持FLOAT32。

    • preTockens：Host侧的int64_t，用于稀疏计算的参数，可选参数，数据类型支持INT64。

    • nextTockens：Host侧的int64_t，用于稀疏计算的参数，可选参数，数据类型支持INT64。

    • headNum：Host侧的int64_t，代表head个数，数据类型支持INT64。

    • inputLayout：Host侧的string，代表输入query、key、value的数据排布格式，支持BSH、SBH、BSND、BNSD。

      :::note 说明 query、key、value数据排布格式支持从多种维度解读，其中B（Batch）表示输入样本批量大小、S（Seq-Length）表示输入样本序列长度、H（Head-Size）表示隐藏层的大小、N（Head-Num）表示多头数、D（Head-Dim）表示隐藏层最小的单元尺寸，且满足D=H/N。 :::

    • innerPrecise：Host侧的int32_t，数据类型支持INT32，内部计算精度模式，其中0表示为高精度，1表示为高性能。

    • workspaceSize：返回用户需要在Device侧申请的workspace大小。

    • executor：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    :::note 说明 第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的query、key、value、dy、dq、dk、dv是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：query、key、value、dy、pseShift、dropMask、paddingMask、attenMask、softmaxMax、softmaxSum、softmaxIn、attentionIn、dq、dk、dv的数据类型和数据格式不在支持的范围内。 :::

aclnnFlashAttentionScoreGrad​

  • 接口定义：

    aclnnStatus aclnnFlashAttentionScoreGrad(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：

    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnFlashAttentionScoreGradGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束与限制​

关于数据shape的约束，以inputLayout的BSND、BNSD为例（BSH、SBH下H=N*D），其中

  • B：取值范围为1~256。
  • N：取值范围为1~256。
  • S：取值范围为1~32K，且为16的倍数。
  • D：取值为64、96、128、256。

父主题： 融合类算子接口

aclnnMoeFFN​

接口原型​

每个算子有两段接口，必须先调用“aclnnMoeFFNGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnMoeFFN”接口执行计算。两段式接口如下：

  • 第一段接口：aclnnStatus aclnnMoeFFNGetWorkspaceSize(const aclTensor *x, const aclIntArray *expertTokens, const aclTensor *weight1, const aclTensor *bias1, const aclTensor *weight2, const aclTensor *bias2, const aclTensor *scale, const aclTensor *offset, const aclTensor *deqScale1, const aclTensor *deqScale2, char *activation, const aclTensor *y, uint64_t *workspaceSize, aclOpExecutor **executor)
  • **第二段接口：****aclnnStatus aclnnMoeFFN(**void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

功能描述​

  • 算子功能：MoE（Mixture-of-Experts, 混合专家系统）是一种用于训练万亿参数量级模型的技术。MoE将预测建模任务分解为若干子任务，在每个子任务上训练一个专家模型（Expert Model），开发一个门控模型（Gating Model），该模型根据要预测的输入来学习信任哪个专家，最终综合多个专家计算结果作为预测结果。

  • 计算公式：

    假设输入样本为x，专家数为N，G(x)代表门控模型函数，FFEi(x)代表第i个专家的输出，所有专家的加权输出为

aclnnMoeFFNGetWorkspaceSize​

  • 接口定义：

    aclnnStatus aclnnMoeFFNGetWorkspaceSize(const aclTensor *x, const aclIntArray *expertTokens, const aclTensor *weight1, const aclTensor *bias1, const aclTensor *weight2, const aclTensor *bias2, const aclTensor *scale, const aclTensor *offset, const aclTensor *deqScale1, const aclTensor *deqScale2, char *activation, const aclTensor *y, uint64_t *workspaceSize, aclOpExecutor **executor)

  • 参数说明：

    • x：Device侧的aclTensor，公式中的输入样本x，数据类型支持FLOAT16、INT8，数据格式支持ND，只支持输入为2维。
    • expertTokens：Host侧的aclIntArray类型，代表各专家的评估权重，数据类型支持INT64，数据格式支持ND，支持的最大长度为256个。
    • weight1：Device侧的aclTensor，专家的评估数据，数据类型支持FLOAT16、INT8，数据格式支持ND，输入为3维。
    • bias1：Device侧的aclTensor，专家的评估数据修正值，可选参数，数据类型支持FLOAT16、INT8，数据格式支持ND。
    • weight2：Device侧的aclTensor，专家的评估数据，数据类型支持FLOAT16、INT8，数据格式支持ND。
    • bias2：Device侧的aclTensor，专家的评估数据修正值，可选参数，数据类型支持FLOAT16、INT8，数据格式支持ND。
    • scale：Device侧的aclTensor，量化参数，可选，数据类型支持FLOAT，数据格式支持ND，一维向量，元素个数与expertTokens长度一致。
    • offset：Device侧的aclTensor，量化参数，可选，数据类型支持FLOAT，数据格式支持ND，一维向量，元素个数与expertTokens长度一致。
    • deqScale1：Device侧的aclTensor，第一个分组matmul的反量化参数，可选，数据类型支持UINT64，数据格式支持ND，输入2维，第一维等于expertTokens长度，第二维等于weight1第三维度。
    • deqScale2：Device侧的aclTensor，第二个分组matmul的反量化参数，可选，数据类型支持UINT64，数据格式支持ND，输入2维，第一维等于expertTokens长度，第二维等于weight2第三维度。
    • activation：Host侧的属性值，代表使用的激活函数，当前仅支持fastgelu。
    • y：Device侧的aclTensor，公式中的输出y，数据类型支持FLOAT、FLOAT16，数据格式支持ND。
    • workspaceSize：返回用户需要在Device侧申请的workspace大小。
    • executor：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    :::note 说明 第一段接口完成入参校验，出现以下场景时报错：

    • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：传入的x、expertTokens、weight1、y或executor是空指针。
    • 返回161002（ACLNN_ERR_PARAM_INVALID）：x、expertTokens、weight1、bias1、weight2、bias2、y的数据类型和数据格式不在支持的范围内。 :::

aclnnMoeFFN​

  • 接口定义：

    aclnnStatus aclnnMoeFFN(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

  • 参数说明：

    • workspace：在Device侧申请的workspace内存起址。
    • workspaceSize：在Device侧申请的workspace大小，由第一段接口aclnnMoeFFNGetWorkspaceSize获取。
    • executor：op执行器，包含了算子计算流程。
    • stream：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束与限制​

  • 专家数据的总数需要与BS（其中B表示输入样本批量大小、S表示输入样本序列长度）长度保持一致；
  • 专家个数不超过256个；
  • 当前MoeFFN实现过程中默认使用fastGelu激活函数，不支持gelu激活函数。

父主题： 融合类算子接口

在线提单