[object Object]

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

• API function: Fuses the query and key operators into one to improve the performance of the inference network. Computes the rotary positional encoding and updates the computation result in place.

  This API has the following function changes based on . Select a proper API based on your actual requirements.

  • The [object Object] parameter is added to control different rotary encoding modes.

• Formula:

  (1) When [object Object] is [object Object]:

  $$query\_q1 = query[..., : query.shape[-1] // 2]$$ $$query\_q2 = query[..., query.shape[-1] // 2 :]$$ $$query\_rotate = torch.cat((-query\_q2, query\_q1), dim=-1)$$ $$key\_k1 = key[..., : key.shape[-1] // 2]$$ $$key\_k2 = key[..., key.shape[-1] // 2 :]$$ $$key\_rotate = torch.cat((-key\_k2, key\_k1), dim=-1)$$ $$q\_embed = (query * cos) + query\_rotate * sin$$ $$k\_embed = (key * cos) + key\_rotate * sin$$

  (2) When [object Object] is [object Object]:

  $$query\_q1 = query[..., : query.shape[-1] // 4]$$ $$query\_q2 = query[..., query.shape[-1] // 4 : query.shape[-1] // 2]$$ $$query\_q3 = query[..., query.shape[-1] // 2 : query.shape[-1] // 4 * 3]$$ $$query\_q4 = query[..., query.shape[-1] // 4 * 3 :]$$ $$query\_rotate = torch.cat((-query\_q2, query\_q1, -query\_q4, query\_q3), dim=-1)$$ $$key\_q1 = key[..., : key.shape[-1] // 4]$$ $$key\_q2 = key[..., key.shape[-1] // 4 : key.shape[-1] // 2]$$ $$key\_q3 = key[..., key.shape[-1] // 2 : key.shape[-1] // 4 * 3]$$ $$key\_q4 = key[..., key.shape[-1] // 4 * 3 :]$$ $$key\_rotate = torch.cat((-key\_q2, key\_q1, -key\_q4, key\_q3), dim=-1)$$ $$q\_embed = (query * cos) + query\_rotate * sin$$ $$k\_embed = (key * cos) + key\_rotate * sin$$

  (3) When [object Object] is [object Object]:

  $$query\_q1 = query[..., ::2].view(-1, 1)$$ $$query\_q2 = query[..., 1::2].view(-1, 1)$$ $$query\_rotate = torch.cat((-query\_q2, query\_q1), dim=-1).view(query.shape[0], query.shape[1], query.shape[2], query.shape[3])$$ $$key\_q1 = key[..., ::2].view(-1, 1)$$ $$key\_q2 = key[..., 1::2].view(-1, 1)$$ $$key\_rotate = torch.cat((-key\_q2, key\_q1), dim=-1).view(key.shape[0], key.shape[1], key.shape[2], key.shape[3])$$ $$q\_embed = (query * cos) + query\_rotate * sin$$ $$k\_embed = (key * cos) + key\_rotate * sin$$

[object Object]

Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters

  [object Object]
  • [object Object]Atlas inference products[object Object]: BFLOAT16 is not supported.

• Returns

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

  The first-phase API implements input parameter verification. The following errors may be thrown.

  [object Object][object Object]

[object Object]

• Parameters

  [object Object]

• Returns

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

[object Object]

• Deterministic computing:

  • [object Object] defaults to a deterministic implementation.

• For [object Object]Atlas inference products[object Object], [object Object]Atlas A3 training products/Atlas A3 inference products[object Object], [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]:

  • When layout is 1, the first two dimensions (B and S) of the input shapes of queryRef, keyRef, cos, and sin must be the same. When layout is 4, the first dimension (T) must be the same.

  • The last dimension (D) of the input shapes of queryRef and keyRef must be the same, and the last dimension (D) of the input shapes of cos and sin must be the same.

  • The dtype of input tensors queryRef, keyRef, cos, and sin must be the same.

  • When layout is 1, the shape of the input queryRef is (q_b, q_s, q_n, q_d), the shape of keyRef is (q_b, q_s, k_n, q_d), and the shape of cos and sin is (q_b, q_s, 1, cos_d). b indicates batch_size, s indicates seq_length, n indicates head_num, and d indicates head_dim. When layout is 4, the shape of the input queryRef is (q_t, q_n, q_d), the shape of keyRef is (q_t, k_n, q_d), and the shape of cos and sin is (q_t, 1, cos_d). t indicates the combined axis of b and s, n indicates head_num, and d indicates head_dim.

    • When the input is BFLOAT16, cast is 1, castSize is 4, and DtypeSize is 2.
    • When the input is FLOAT16 or FLOAT32, cast is 0, and castSize = DtypeSize (2 for FLOAT16 and 4 for FLOAT32).

    lastDim indicates the dimension of the input shape that requires positional encoding to be rotated. The size of the required UB space is calculated as follows: ub_required = (q_n + k_n) lastDim castSize 2 + lastDim DtypeSize 4 + (q_n + k_n) lastDim castSize + (q_n + k_n) lastDim castSize 2 + cast (lastDim 4 * 2), If the value of ub_required exceeds the total UB space of the current AI processor, this fusion operator cannot be used.

• Ascend 950PR/Ascend 950DT:

  • For any layout, the dimensions of queryRef and keyRef must be the same except the N dimension. The last dimension (D) of the input shapes of queryRef and keyRef must be the same, and the last dimension (D) of the input shapes of cos and sin must be the same and less than or equal to the last dimension (D) of the input shapes of queryRef and keyRef.
  • The dtype of input tensors queryRef, keyRef, cos, and sin must be the same.
  • When rotaryMode is set to "half" or "interleave", the last dimension of the input shape must be exactly divided by 2. When rotaryMode is set to "quarter", the last dimension of the input shape must be exactly divided by 4.

[object Object]

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