[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.

• Formulas:

  1. In mrope mode, the shape of positions is [m, numTokens], where m indicates the number of elements in mropeSection, which can be 3 or 4.
  $$cosSin[i] = cosSinCache[positions[i]]$$ $$cos, sin = cosSin.chunk(2, dim=-1)$$

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

  (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, aclnnRopeWithSinCosCacheGetWorkspaceSize is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, aclnnRopeWithSinCosCache 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:
  • [object Object] defaults to a deterministic implementation.
• 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 is a multiple of 32. If the data type is FLOAT32, the value is a multiple of 16.
• rotaryDim: The value is always less than or equal to headSize. If the data type is BFLOAT16 or FLOAT16, the value is a multiple of 32. If the data type is FLOAT32, the value is a multiple of 16. In mrope mode, the sum of all elements in mropeSection must be half of the rotaryDim value.
• The value of the input tensor positions must be less than the 0D maxSeqLen of cosSinCache.
• In mrope mode, mropeSection can only be [16, 24, 24], [24, 20, 20], [8, 12, 12], or [16, 16, 16, 16].

[object Object]

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