RopeOperation（代码开放）

产品支持情况

功能说明

旋转位置编码（Rotary Position Embedding，RoPE），以旋转矩阵的方式在q、k中注入位置信息，使得attention计算时能感受到token的位置关系，在各大模型中，RoPE被广泛应用。RoPE以绝对位置编码的方式实现了相对位置编码，能有效保持位置信息相对关系，并且可以通过编码外推的方式支持超过训练长度的位置编码。

算子上下文

图1 RopeOperation算子上下文

计算公式

对于二维情况

假设空间是偶数维的，把原始空间切分一个个正交的二维子空间，在上面做独立的不同角度的旋转，可以扩展到高维空间。

定义

1
2
3
4
5

struct RopeParam {
    int32_t rotaryCoeff = 4;
    int32_t cosFormat = 0;
    uint8_t rsv[8] = {0};
};

参数列表

rotaryCoeff参数选择与原始计算公式的对应关系如下：

其中m为token的位置，d为query或key的维度。

• rotaryCoeff = 2的情况：

  如上式，query和key分别取一半，进行旋转操作后拼接，称为对半旋转，对应rotaryCoeff = 2，即1/2。

• rotaryCoeff = 4的情况：

  以此类推，rotaryCoeff = 4对应于1/4，表示query和key分成前后两半，每一半按rotaryCoeff = 2的情况处理。

• rotaryCoeff = head_size的情况：

  如上式，query和key中的维度两两一组交错排列，其中图示的cos和sin列表经过了concat操作，表现为连续相同的2个cos、sin值。

  如果不经过concat操作，cos和sin列表的长度为上式中的一半，即head_size/2。算子也支持此种情况，此时需要将rotaryCoeff = head_size/2。

  换言之，使用上图的计算模式只需要满足rotaryCoeff与cos或sin输入tensor的最后一维大小相等。

输入

输出

约束说明

• 输入tensor数据类型需保持一致，高精度模式例外。
• cos、sin传入数据类型为float时，中间计算结果以float保存。
• hiddenSizeQ和hiddenSizeK必须是head_size的整数倍，满足hiddenSizeQ = head_size * headNumQ、hiddenSizeK = head_size * headNumK，其中headNumQ可以大于headNumK，hiddenSizeQ和hiddenSizeK需要32bytes对齐。
• ntokens = sum(seqlen[i])，i=0...batch-1。
• query和key要求两维，有部分模型使用了4维，这种情况下维度是：

  [batch, seqlen, headNum, head_size]；对应的ropeQ、ropeK也是四维，维度输入输出对应。

• Decoder阶段要取cos和sin表中seqlen对应的cos/sin值输入。
• 多batch场景需要组合使用gather算子。