RopeOperation

Description

Rotary Position Embedding (RoPE): Location information is injected into q and k in a rotation matrix, so that a location relationship of a token can be sensed during attention calculation. RoPE is widely used in various models. RoPE implements relative position coding by using absolute position coding, can effectively maintain a relative relationship between position information, and can support position coding that exceeds a training length by using coding extrapolation.

Operator Context

Formula

For a two-dimensional case:

Assuming that the space is of an even dimension, the original space is divided into orthogonal two-dimensional subspaces, and independent rotation at different angles is performed on the subspaces, so that the space can be extended to a high-dimensional space.

Definition

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

Parameters

The mapping between the rotaryCoeff parameter and the original formula is as follows:

m is the location of the token, and d is the dimension of the query or key.

• When rotaryCoeff = 2:

  As shown in the preceding formula, take half of the query and key values, perform the rotation operation, and concatenate the values. This is called half rotation, which corresponds to rotaryCoeff = 2, that is, 1/2.

• When rotaryCoeff = 4:

  Similarly, rotaryCoeff = 4 corresponds to 1/4, indicating that the query and key are divided into two parts, and each half is processed based on rotaryCoeff = 2.

• When rotaryCoeff = headDim:

  As shown in the preceding formula, the dimensions in the query and key are arranged in a staggered manner. The cos and sin lists in the figure are concatenated and represented as two consecutive same cos and sin values.

  If the concat operation is not performed, the length of the cos and sin lists is half of that in the preceding formula, that is, headDim/2. The operator also supports this case. In this case, rotaryCoeff = headDim/2.

  In other words, to use the calculation mode in the preceding figure, only need to ensure that the size of rotaryCoeff is the same as that of the last dimension of the input tensor of cos or sin.

Input

Output

Restrictions

• The data types of the input tensors must be the same, except for the high-precision mode.
• If the input data type of cos and sin is float, the intermediate calculation result is saved in float format.
• The values of hiddenSizeQ and hiddenSizeK must be integer multiples of the value of headDim and meet the requirements of hiddenSizeQ = headDim * headNumQ and hiddenSizeK = headDim * headNumK. The value of headNumQ can be greater than that of headNumK.
• ntokens = sum(seqlen[i]), i=0...batch-1.
• The query and key must be two-dimensional. Some models use four dimensions. In this case, the dimensions are as follows:

  [batch, seqlen, headNum, headDim]. The corresponding ropeQ and ropeK are also four-dimensional, and the input and output dimensions correspond to each other.

• In the decoder phase, the cos/sin value corresponding to seqlen in the cos and sin tables is used.
• In the multi-batch scenario, the gather operator needs to be used together.