昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

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两路算子融合成一路。
  • 计算公式：queryq1=query[...,:query.shape[−1]//2]query_q1 = query[..., : query.shape[-1] // 2] queryq2=query[...,query.shape[−1]//2:]query_q2 = query[..., query.shape[-1] // 2 :] queryrotate=torch.cat((−queryq2,queryq1),dim=−1)query_rotate = torch.cat((-query_q2, query_q1), dim=-1) keyk1=key[...,:key.shape[−1]//2]key_k1 = key[..., : key.shape[-1] // 2] keyk2=key[...,key.shape[−1]//2:]key_k2 = key[..., key.shape[-1] // 2 :] keyrotate=torch.cat((−keyk2,keyk1),dim=−1)key_rotate = torch.cat((-key_k2, key_k1), dim=-1) qembed=(query∗cos)+queryrotate∗sinq_embed = (query * cos) + query_rotate * sin kembed=(key∗cos)+keyrotate∗sink_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。

[object Object]