aclnnRopeWithSinCosCache

产品支持情况

产品	是否支持
[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]	×

功能说明

• 算子功能：推理网络为了提升性能，将sin和cos输入通过cache传入，执行旋转位置编码计算。

• 计算公式：

  1、mrope模式：positions的shape输入是[3, numTokens]：

  $$cosSin[i] = cosSinCache[positions[i]]$$ $$cos, sin = cosSin.chunk(2, dim=-1)$$ $$cos0 = cos[0, :, :mropeSection[0]]$$ $$cos1 = cos[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])]$$ $$cos2 = cos[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])]$$ $$cos = torch.cat((cos0, cos1, cos2), dim=-1)$$ $$sin0 = sin[0, :, :mropeSection[0]]$$ $$sin1 = sin[1, :, mropeSection[0]:(mropeSection[0] + mropeSection[1])]$$ $$sin2 = sin[2, :, (mropeSection[0] + mropeSection[1]):(mropeSection[0] + mropeSection[1] + mropeSection[2])]$$ $$sin= torch.cat((sin0, sin1, sin2), dim=-1)$$ $$queryRot = query[..., :rotaryDim]$$ $$queryPass = query[..., rotaryDim:]$$

  （1）rotate_half（GPT-NeoX style）计算模式：

  $$x1, x2 = torch.chunk(queryRot, 2, dim=-1)$$ $$o1[i] = x1[i] * cos[i] - x2[i] * sin[i]$$ $$o2[i] = x2[i] * cos[i] + x1[i] * sin[i]$$ $$queryRot = torch.cat((o1, o2), dim=-1)$$ $$query = torch.cat((queryRot, queryPass), dim=-1)$$

  （2）rotate_interleaved（GPT-J style）计算模式：

  $$x1 = queryRot[..., ::2]$$ $$x2 = queryRot[..., 1::2]$$ $$queryRot = torch.stack((o1, o2), dim=-1)$$ $$query = torch.cat((queryRot, queryPass), dim=-1)$$

  2、rope模式：positions的shape输入是[numTokens]：

  $$cosSin[i] = cosSinCache[positions[i]]$$ $$cos, sin = cosSin.chunk(2, dim=-1)$$ $$queryRot = query[..., :rotaryDim]$$ $$queryPass = query[..., rotaryDim:]$$

  （1）rotate_half（GPT-NeoX style）计算模式：

  $$x1, x2 = torch.chunk(queryRot, 2, dim=-1)$$ $$o1[i] = x1[i] * cos[i] - x2[i] * sin[i]$$ $$o2[i] = x2[i] * cos[i] + x1[i] * sin[i]$$ $$queryRot = torch.cat((o1, o2), dim=-1)$$ $$query = torch.cat((queryRot, queryPass), dim=-1)$$

  （2）rotate_interleaved（GPT-J style）计算模式：

  $$x1 = query\_rot[..., ::2]$$ $$x2 = query\_rot[..., 1::2]$$ $$queryRot = torch.stack((o1, o2), dim=-1)$$ $$query = torch.cat((queryRot, queryPass), dim=-1)$$

函数原型

每个算子分为undefined，必须先调用“aclnnRopeWithSinCosCacheGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnRopeWithSinCosCache”接口执行计算。

• aclnnStatus aclnnRopeWithSinCosCacheGetWorkspaceSize(const aclTensor *positions, const aclTensor *queryIn, const aclTensor *keyIn, const aclTensor *cosSinCache, const aclIntArray *mropeSection, int64_t headSize, bool isNeoxStyle, aclTensor *queryOut, aclTensor *keyOut, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnRopeWithSinCosCache(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnRopeWithSinCosCacheGetWorkspaceSize

• 参数说明：

  • positions（aclTensor*，计算输入）：Device侧的aclTensor，输入索引，公式中的positions，用于选取位置编码张量。要求是一个维度为1D或2D的Tensor，shape为(numTokens)或(3, numTokens)，1D维度输入是rope模式，2D维度输入是mrope模式。numTokens表示一个序列中的token数量。支持undefined，支持空Tensor。mrope/rope模式下数据类型支持INT32、INT64，undefined支持ND。
  • queryIn（aclTensor*，计算输入）：Device侧的aclTensor，表示要执行旋转位置编码的第一个张量，公式中的query，要求是一个维度为2D的Tensor，shape为(numTokens, numQHeads*headSize)。numQHeads表示query的注意力头数量。headSize表示每个注意力头维度大小。支持undefined，支持空Tensor。mrope/rope模式下数据类型支持BFLOAT16、FLOAT16、FLOAT32，undefined支持ND。
  • keyIn（aclTensor*，计算输入）：Device侧的aclTensor，表示要执行旋转位置编码的第二个张量，公式中的key，要求是一个维度为2D的Tensor，shape为(numTokens, numKHeads*headSize)。numKHeads表示key的注意力头数量。headSize表示每个注意力头维度大小。支持undefined，支持空Tensor。mrope/rope模式下数据类型支持BFLOAT16、FLOAT16、FLOAT32，undefined支持ND。
  • cosSinCache（aclTensor*，计算输入）：Device侧的aclTensor，表示参与计算的位置编码张量，要求shape为一个2D的(maxSeqLen, rotaryDim)。maxSeqLen表示模型处理的序列的最大长度。rotaryDim表示旋转位置嵌入的维度大小。支持undefined，支持空Tensor。mrope/rope模式下数据类型支持BFLOAT16、FLOAT16、FLOAT32，undefined支持ND。
  • mropeSection（aclIntArray*，计算输入）：mrope模式下用于整合输入的位置编码张量信息，公式中的mropeSection，输入mropeSection属性表示使能mrope模式。不使能mrope模式（即rope模式）输入为nullptr。
  • headSize（int64_t, 计算输入）：表示每个注意力头维度大小。数据类型int64。
  • isNeoxStyle（bool, 计算输入）：true表示rotate_half（GPT-NeoX style）计算模式，false表示rotate_interleaved（GPT-J style）计算模式。
  • queryOut（aclTensor*，计算输出）：输出query执行旋转位置编码后的结果，要求是一个2D的Tensor，shape为(numTokens, numQHeads*headSize)。数据类型同query，mrope/rope模式下支持FLOAT、FLOAT16、BFLOAT16，undefined要求为ND。支持输出连续的Tensor。
  • keyOut（aclTensor*，计算输出）：输出key执行旋转位置编码后的结果，要求是一个2D的Tensor，shape为(numTokens, numKvHeads*headSize)。数据类型同key，mrope/rope模式下支持FLOAT、FLOAT16、BFLOAT16，undefined要求为ND。输出连续的Tensor。
  • workspaceSize（uint64_t*，出参）：返回用户需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnRopeWithSinCosCache

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

• 返回值

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

约束说明

• queryIn、keyIn、cosSinCache只支持2维shape输入。
• numQHeads支持范围: 1~32。
• numKHeads支持范围: 1~32。
• headSize支持范围: 16~128。数据类型为BFLOAT16或FLOAT16时为32的倍数，数据类型为FLOAT32时为16的倍数；
• rotaryDim支持范围: 16~128，始终小于等于headSize。数据类型为BFLOAT16或FLOAT16时为32的倍数，数据类型为FLOAT32时为16的倍数；
• 当输入tensor positions中值域超过cosSinCache的0维maxSeqLen，会有越界报错。

调用示例

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