DynamicNTKOperation（代码开放）

产品支持情况

功能说明

推理的长度大于训练序列长度时，生成ROPE算子使用的旋转矩阵（即下面公式中的sin/cos矩阵）。

推理长度小于等于训练序列长度时，不进行插值；推理长度大于训练长度时，放大base动态插值。

将输入的token序列的位置信息positionIds和invFreq进行外积，再cos/sin运算得到最终的Rotary Embedding的结果。

算子上下文

图1 DynamicNTK算子在模型中的位置

DynamicNTK算子输出sin，cos。

q，k，sin，cos一起输入到Rope算子中输出旋转后的q，k：

算子功能分析

DynamicNTK用于推理长度大于训练长度时，对embedding做特殊处理，根据seqLen长度，对base参数进行缩放，最终生成新的cos、sin。由于Llama、Qwen、Yarn等模型对于cos、sin的生成方案存在一部分差异，DynamicNTK算子实现不同模型下公共的处理逻辑，差异化处理逻辑由模型实现。

• 算子计算公式

  将输入的token序列的位置信息positionIds和invFreq进行外积，再cos/sin运算得到最终的Rotary Embedding的结果。

• 具体计算过程：

  先通过每个batch的seqLen取出每个batch的positionIds，invFreq

  对每一个batch：

  freq = torch.einsum('i,j->ij', positionIds, invFreq)
  emb = torch.cat((freq, freq), dim=-1)
  sinOut = emb.sin()
  cosOut = emb.cos()

定义

1
2
3
4

struct DynamicNTKParam {
    aclDataType outDataType = ACL_DT_UNDEFINED;
    uint8_t rsv[12] = {0};
};

参数列表

输入

参数InvFreqs的计算公式为：

输出

约束说明

• headSize（即输出tensor的最后一维）小于等于2048，并且是32的倍数。
• batch小于等于16。
• nTokens（即输入positionIds的维度）小于等于256000。
• InvFreqIn数组数据范围：[0,1)，seqLens数组中数据大于0，且数组和为nTokens。