- 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.
- Formula:
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.
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]
Deterministic computing: The aclnnApplyRotaryPosEmb is implemented in deterministic mode by default.
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 dimensions (D) of the input shapes of queryRef, keyRef, 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 queryRef is represented by (q_b, q_s, q_n, q_d), the shape of keyRef is represented by (q_b, q_s, k_n, q_d), and the shape of cos and sin 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 layout is 4, the shape of queryRef is represented by (q_t, q_n, q_d), the shape of keyRef is represented by (q_t, k_n, q_d), and the shape of cos and sin 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:
[object Object], If the value of[object Object]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 shape of queryRef and keyRef must be the same, and the last dimension (D) of the input shape of cos and sin must be the same and less than or equal to the last dimension (D) of the input shape of queryRef and keyRef.
- The dtype of input tensors queryRef, keyRef, cos, and sin must be the same.
The following example is for reference only. For details, see .