aclnnApplyRotaryPosEmb

支持的产品型号

• Atlas 推理系列产品。
• Atlas A2训练系列产品/Atlas 800I A2推理产品。

接口原型

每个算子分为undefined，必须先调用“aclnnApplyRotaryPosEmbGetWorkspaceSize”接口获取入参并根据流程计算所需workspace大小，再调用“aclnnApplyRotaryPosEmb”接口执行计算。

• aclnnStatus aclnnApplyRotaryPosEmbGetWorkspaceSize(aclTensor *queryRef, aclTensor *keyRef, const aclTensor *cos, const aclTensor *sin, int64_t layoutOptional, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnApplyRotaryPosEmb(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor， aclrtStream stream)

功能描述

• 算子功能：推理网络为了提升性能，将query和key两路算子融合成一路。
• 计算公式：

$$query\_q1 = query[..., : query.shape[-1] // 2]$$ $$query\_q2 = query[..., query.shape[-1] // 2 :]$$ $$query\_rotate = torch.cat((-query\_q2, query\_q1), dim=-1)$$ $$key\_k1 = key[..., : key.shape[-1] // 2]$$ $$key\_k2 = key[..., key.shape[-1] // 2 :]$$ $$key\_rotate = torch.cat((-key\_k2, key\_k1), dim=-1)$$ $$q\_embed = (query * cos) + query\_rotate * sin$$ $$k\_embed = (key * cos) + key\_rotate * sin$$

aclnnApplyRotaryPosEmbGetWorkspaceSize

• 参数说明：

  • queryRef(aclTensor*, 计算输入)：输入Tensor，数据类型支持BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持）、FLOAT16、FLOAT32。支持非连续的Tensor，undefined支持ND。
  • keyRef(aclTensor*, 计算输入)：输入Tensor，数据类型支持BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持）、FLOAT16、FLOAT32。支持非连续的Tensor，undefined支持ND。
  • cos(const aclTensor*, 计算输入)：输入Tensor，数据类型支持BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持）、FLOAT16、FLOAT32。支持非连续的Tensor，undefined支持ND。
  • sin(const aclTensor*, 计算输入)：输入Tensor，数据类型支持BFLOAT16（仅Atlas A2训练系列产品/Atlas 800I A2推理产品支持）、FLOAT16、FLOAT32。支持非连续的Tensor，undefined支持ND。
  • layoutOptional(int64_t, 计算输入)：公式中的布局格式，数据类型int64, 目前只支持1, 代表"BSH"格式，实际代表支持格式为BSND，4维。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

• 返回值：

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

• 说明： 第一段接口完成入参校验，若出现以下错误码，则对应原因为：

  • 返回161001（ACLNN_ERR_PARAM_NULLPTR）：如果传入参数是必选输入，输出或者必选属性，且是空指针，则返回161001。
  • 返回161002（ACLNN_ERR_PARAM_INVALID）：queryRef、keyRef、cos、sin、layoutOptional的数据类型和数据格式不在支持的范围内。

aclnnApplyRotaryPosEmb

• 参数说明：

  • workspace(void*, 入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnApplyRotaryPosEmbGetWorkspaceSize获取。
  • executor(aclOpExecutor*, 入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream, 入参)：指定执行任务的AscendCL Stream流。

• 返回值：

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

约束与限制

• 输入张量queryRef、keyRef、cos、sin只支持4维的shape，layoutOptional只支持1
• 输入张量queryRef、keyRef、cos、sin 的dtype必须相同，且4个输入shape的前2维和最后一维必须相等，cos和sin的shape第3维必须等于1，输入shape最后一维必须等于128
• 输入queryRef的shape用（q_b,q_s,q_n,q_d）表示，keyRef shape用（q_b,q_s,k_n,q_d）,cos和sin shape用（q_b,q_s,1,q_d），当输入是BFLOAT16时，cast表示为1，castSize为4， DtypeSize为2，不是BFLOAT16时，cast表示为0，castSize = DtypeSize，float16时为2，float32时为4， 需要使用的ub大小为：ub = (q_n + k_n)128 castSize2 + 128DtypeSize4 + (q_n + k_n)128castSize + (q_n + k_n)128castSize2 + cast*(12842) 当ub大于对应可以获取到的AI处理器版本总ub大小时，不支持。
• 不支持空tensor场景

调用示例

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