DynamicNTKOperation

Description

If the inference length is greater than the training length, special processing is required for embedding.

If the inference length is less than or equal to the training length, no interpolation is performed. If the inference length is greater than the training length, the base dynamic interpolation is amplified.

An outer product is performed on the position information positionIds and inv_freq of the input token sequence, and then a cos/sin operation is performed to obtain a final result of rotary embedding.

Operator Context

Figure 1 Location of the DynamicNTK operator in the model

The dynamicNTK operator outputs sin and cos.

Input q, k, sin, and cos to the Rope operator and output the rotated q and k.

Operator Function Analysis

DynamicNTK is used to perform special processing on embedding when the inference length is greater than the training length. The base parameter is scaled based on the seqlen to generate new cos and sin. Models such as Llama, Qwen, and Yarn have different cos and sin generation solutions. The DynamicNTK operator implements the common processing logic of different models, and the differentiated processing logic is implemented by the models.

• Operator calculation formula

  An outer product is performed on the position information positionIds and inv_freq of the input token sequence, and then a cos/sin operation is performed to obtain a final result of rotary embedding.

• Details:

  Obtain positionIds and invFreq of each batch based on seqlen of each batch.

  For each batch:

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

Definition

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

Parameters

Input

The formula for calculating InvFreqs is as follows:

Output

Restrictions

• headDim (the last dimension of the output tensor) is less than or equal to 2048 and is a multiple of 32.
• The value of batch is less than or equal to 16.
• The value of ntokens (dimension of positionIds) is less than or equal to 256000.
• Value range of the InvFreqIn array: [0, 1). The value in the seqlens array is greater than 0, and the sum of the arrays is ntokens.