aclnnRotaryPositionEmbedding

产品支持情况

[object Object]undefined

功能说明

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

• 计算公式：

  • [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]、[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[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”接口执行计算。

• [object Object]
• [object Object]

aclnnRotaryPositionEmbeddingGetWorkspaceSize

• 参数说明：

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

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

  • sin（aclTensor*，计算输入）：表示参与计算的位置编码张量，公式中的[object Object]，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 A2 推理系列产品[object Object]：
      • 0：表示旋转模式为half模式。
      • 1：表示旋转模式为interleave模式。

  • out（aclTensor*，计算输出）：表示旋转位置编码的计算结果张量，公式中的[object Object]，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 A2 推理系列产品[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
      • 当（D/2）% (32/inputDtypeSize) == 0时，需满足B * N <= S * 8
      • 当（D/2）% (32/inputDtypeSize) != 0时，需满足B * N * 2 <= (S + coreNum -1) / coreNum 或者 D >= 80
    • 当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

调用示例

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