aclnnRotaryPositionEmbedding

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：执行单路旋转位置编码计算。

• 计算公式：

  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：

  （1）half模式（mode等于0）： $$ x1 = x[..., : x.shape[-1] // 2] $$

  $$x2 = x[..., x.shape[-1] // 2 :]$$ $$x\_rotate = torch.cat((-x2, x1), dim=-1)$$ $$y = x * cos + x\_rotate * sin$$
  [object Object]
  $$x1 = x[..., ::2].view(-1, 1)$$ $$ x2 = x[..., 1::2].view(-1, 1) $$ $$

  x_rotate = torch.cat((-x2, x1), dim=-1).view(x.shape[0], x.shape[1], x.shape[2], x.shape[3]) $$
  $$ y = x * cos + x_rotate * sin $$

  （3）quarter模式（mode等于2）： $$ x1 = x[..., : x.shape[-1] // 4] $$

  $$ x2 = x[..., x.shape[-1] // 4 : x.shape[-1] // 2] $$ $$

  x3 = x[..., x.shape[-1] // 2 : x.shape[-1] // 4 * 3] $$
  $$ x4 = x[..., x.shape[-1] // 4 * 3 :] $$

  $$x\_rotate = torch.cat((-x2, x1, -x4, x3), dim=-1)$$ $$ y = x * cos + x\_rotate * sin $$ （4）interleave-half模式（mode等于3），该模式会先将奇数位的输入抽取到前半部分，将偶数位的输入抽取到后半部分，再进行half处理： $$

  x1 = x[..., ::2] $$

  $$ x2 = x[..., 1::2] $$ $$

  x_part1 = torch.cat((x1, x2), dim=-1) $$

  $$ x\_part2 = torch.cat((-x2, x1), dim=-1) $$ $$

  y = x_part1 * cos + x_part2 * sin $$

函数原型

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

• aclnnStatus aclnnRotaryPositionEmbeddingGetWorkspaceSize(const aclTensor *x, const aclTensor *cos, const aclTensor *sin, int64_t mode, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnRotaryPositionEmbedding(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnRotaryPositionEmbeddingGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：表示待执行旋转位置编码的张量，公式中的x，Device侧的aclTensor。支持，支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32。

  • cos（aclTensor*，计算输入）：表示参与计算的位置编码张量，公式中的cos，Device侧的aclTensor。支持，支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。

  • sin（aclTensor*，计算输入）：表示参与计算的位置编码张量，公式中的sin，Device侧的aclTensor，支持，支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。

  • mode（int64_t，计算输入）：表示公式中的旋转模式，数据类型为int64，支持的枚举值如下：

    • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：
      • 0：表示旋转模式为half模式。
      • 1：表示旋转模式为interleave模式。

  • out（aclTensor*，计算输出）：表示旋转位置编码的计算结果张量，公式中的y，Device侧的aclTensor。 不支持，支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。

  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。

  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnRotaryPositionEmbedding

• 参数说明：

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

• 返回值：

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

约束说明

• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]：

  输入张量x支持BNSD、BSND、SBND排布。 输入张量x、cos、sin及输出张量y的D维度大小必须相同，满足D<896，且必须为2的倍数。 输入张量x和输出张量y的shape必须完全相同。 输入张量cos和sin的shape必须完全相同.

  • half模式：
    • B，N < 1000;
    • 当x为BNSD时，cos、sin支持11SD、B1SD、BNSD
    • 当x为BSND时，cos、sin支持1S1D、BS1D、BSND
    • 当x为SBND时，cos、sin支持S11D、SB1D、SBND
  • interleave模式：
    • B * N < 1000
    • 当x为BNSD时，cos、sin支持11SD
    • 当x为BSND时，cos、sin支持1S1D
    • 当x为SBND时，cos、sin支持S11D

调用示例

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