[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]: The BFLOAT16 data type is not supported.

• Returns:

  aclnnStatus: status code. For details, see .

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

  [object Object]

[object Object]

• Parameters

  [object Object]

• Returns:

  aclnnStatus: 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 A2 training products/Atlas A2 inference products[object Object], [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]:

    • When [object Object] is 1, the first two dimensions (B and S) of the input shapes of [object Object], [object Object], [object Object], and [object Object] must be the same. When [object Object] is 4, the first dimension (T) must be the same.

    • The last dimensions (D) of the input shapes of [object Object], [object Object], [object Object], and [object Object] must be the same.

    • The dtype of input tensors [object Object], [object Object], [object Object], and [object Object] must be the same.

    • When [object Object] is 1, the shape of [object Object] is represented by (q_b, q_s, q_n, q_d), the shape of [object Object] is represented by (q_b, q_s, k_n, q_d), and the shape of [object Object] and [object Object] is represented by (q_b, q_s, 1, q_d). b indicates batch_size, s indicates seq_length, n indicates head_num, and d indicates head_dim. When [object Object] is 4, the shape of [object Object] is represented by (q_t, q_n, q_d), the shape of [object Object] is represented by (q_t, k_n, q_d), and the shape of [object Object] and [object Object] is represented by (q_t, 1, q_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 value of head_dim in the last dimension of the input shape. The UB space size to be used 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.

[object Object]

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