[object Object]

[object Object][object Object]undefined[object Object]

• Description: For the inference network, the sin and cos inputs are transferred through the cache to perform rotation position encoding computation. Compared with the API, this API has the cacheMode parameter added to indicate the mode of concatenating cos and sin.

  • When cacheMode is set to 0, the implementation is the same as that of , that is, cos and sin are concatenated in segment mode.
  • When cacheMode is 1, cos and sin are staggered.

• Formulas:

  1. MRoPE mode: The shape of [object Object] is [object Object], where [object Object] represents the number of elements in [object Object]. Valid values are [object Object] or [object Object].
  $$cosSin[i] = cosSinCache[positions[i]]$$ $$cos, sin = cosSin.chunk(2, dim=-1)$$

  (1) When cacheMode is set to 0:

  • When [object Object] contains 3 elements:

    $$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:]$$

  • When [object Object] contains 4 elements:

    $$cos = torch.cat([m[i]\ for\ i, m\ in\ enumerate(cos.split(mropeSection, dim=-1))], dim=-1)$$ $$sin = torch.cat([m[i]\ for\ i, m\ in\ enumerate(sin.split(mropeSection, dim=-1))], dim=-1)$$ $$queryRot = query[..., :rotaryDim]$$ $$queryPass = query[..., rotaryDim:]$$

  (2) When cacheMode is set to 1:

  $$cosTmp = cos$$ $$cos [..., 1:mropeSection[1] * 3:3] = cosTmp[1, ..., 1:mropeSection[1] * 3:3]$$ $$cos[..., 2:mropeSection[1] * 3:3] = cosTmp[2, ..., 2:mropeSection[1] * 3:3]$$ $$sinTmp = sin$$ $$sin[..., 1:mropeSection[1] * 3:3] = sinTmp [1, ..., 1:mropeSection[1] * 3:3]$$ $$sin[..., 2:mropeSection[1] * 3:3] = sinTmp [2, ..., 2:mropeSection[1] * 3:3]$$ $$queryRot = query[..., :rotaryDim]$$ $$queryPass = query[..., rotaryDim:]$$

  (1) rotate_half (GPT-NeoX style) computation mode:

  $$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) computation mode:

  $$x1 = queryRot[..., ::2]$$ $$x2 = queryRot[..., 1::2]$$ $$o1[i] = x1[i] * cos[i] - x2[i] * sin[i]$$ $$o2[i] = x2[i] * cos[i] + x1[i] * sin[i]$$ $$queryRot = torch.stack((o1, o2), dim=-1)$$ $$query = torch.cat((queryRot, queryPass), dim=-1)$$
  2. RoPE mode: The shape of positions is [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) computation mode:

  $$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) computation mode:

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

[object Object]

Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters:

  [object Object]

• Returns

  [object Object]: status code. For details, see .

  The first-phase API implements input parameter validation. The following error codes may be returned.

  [object Object]

[object Object]

• Parameters:

  [object Object]

• Returns

  [object Object]: status code. For details, see .

[object Object]

• Deterministic computation:
  • The aclnnRopeWithSinCosCacheV2 function is implemented in deterministic mode by default.
• queryIn, keyIn, and cosSinCache support only 2D shape input.
• The data types of queryIn, keyIn, and cosSinCache must be the same.
• headSize: If the data type is BFLOAT16 or FLOAT16, the value must be a multiple of 32. If the data type is FLOAT32, the value must be a multiple of 16.
• rotaryDim: The value must be less than or equal to headSize. If the data type is BFLOAT16 or FLOAT16, the value must be a multiple of 32. If the data type is FLOAT32, the value must be a multiple of 16. In mrope mode, the sum of all elements in mropeSection must be half of the value of rotaryDim.
• The value of the input tensor positions must be less than the value of maxSeqLen in dimension 0 of cosSinCache.
• In mrope mode, mropeSection can only be [16, 24, 24], [24, 20, 20], [8, 12, 12], or [16, 16, 16, 16].
• In mrope mode, cacheMode can only be 0 or 1. When mropeSection is set to [16, 16, 16, 16], cacheMode can only be 0.

[object Object]

The following example is for reference only. For details, see .