Uses the FlashAttention algorithm to perform self-attention computation in training scenarios. The query, key, and value parameters of this API support multiple sequences with the same length or different lengths.
- Compared with the API, this API has the following differences:
- Adjusting the Dropout function: When keepProb is less than 1.0, if no external DropoutMask is passed, the new parameters seed and offset are used to generate the DropoutMask. If an external DropoutMask is passed, the external DropoutMask is used.
- Compared with the API, this API has the following differences:
- Adjusting the Dropout function: When keepProb is less than 1.0, if no external DropoutMask is passed, the new parameters seed and offset are used to generate the DropoutMask. If an external DropoutMask is passed, the external DropoutMask is used.
- Compared with the API, this API has the following differences:
Formulas:
The forward propagation formula for attention is as follows:
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.
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]
Deterministic computation:
- The default deterministic implementation of aclnnFlashAttentionScoreV4.
When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
Input query, key, and value
- B: The batch sizes must be equal.
- D: Head-Dim must satisfy (qD == kD && kD >= vD).
[object Object]must be consistent.
The input shape of queryRopeOptional is the same as that of query in all dimensions except the D dimension.
The input shape of keyRopeOptional is the same as that of key in all dimensions except the D dimension.
The restrictions on the data shape are as follows, using TND, BSND, and BNSD of inputLayout as examples (H = N x D in BSH and SBH):
- T(B x S): The value ranges from 1 to 1M. In TND format, the maximum length supported by actualSeqQLenOptional is 20,000.
[object Object]: The value ranges from 1 to 2M. When[object Object]is passed, B supports a maximum of 2K.- N: The value ranges from 1 to 256.
- S: The value ranges from 1 to 1M.
- D: The value ranges from 1 to 768.
The data layout of query, key, and value can be interpreted from multiple dimensions. To be specific, B (Batch) indicates the size of an input sample batch, S (Seq-Length) indicates the length of the input sample sequence, H (Head-Size) indicates the size of the hidden layer, N (Head-Num) indicates the number of heads, and D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N).
[object Object]: 0 and 1 are reserved, and 2 indicates that invalid row calculation is enabled. This function is used to prevent precision loss caused by the mask of the entire row during calculation. However, this configuration deteriorates the performance. If the operator can determine that invalid rows exist, the invalid row computation is automatically enabled, such as in scenarios where[object Object]is set to 3 and Sq is greater than Skv.Meanings of
[object Object]values:[object Object]undefined
When
[object Object]is set to 2 or 3, Sq and Skv must be of the same length.The constraints on sparseMode are as follows:
- When the shapes of all
[object Object]are less than 2048 and are the same, the default mode is recommended to reduce memory usage. - When the value is set to 1, 2, 3, or 5, the user-configured
[object Object]and[object Object]do not take effect. - When the value is set to 0 or 4, ensure that the ranges of
[object Object],[object Object], and[object Object]are consistent. - If no specific value is required, 0 is recommended.
- For details about the sparse modes, see .
- When the shapes of all
In some scenarios, if the computation load is too large, the operator execution may time out (an AI Core error is reported, and
[object Object]is[object Object]). In this case, you are advised to perform axis splitting. Note: The computation load is affected by parameters such as B, S, N, and D. Larger values indicate larger computation loads.In the band scenario, the values of
[object Object]and[object Object]must overlap.In the prefixOptional sparse computing scenario, sparseMode is set to 5 or 6 when the sequence lengths are the same, and sparseMode is set to 6 when the sequence lengths are different. In the two scenarios, when Sq > Skv, the value range of N for prefix is [0, Skv]. When Sq <= Skv, the value range of N for prefix is [Skv-Sq, Skv]. If sparseModeOptional is set to 5 and prefix N > Skv or prefixOptional is not passed, full computation is performed. If sparseModeOptional is set to 6, prefixOptional must be passed.
If Sq of
[object Object]is greater than 1024, if BNHS and 1NHS are configured, Sq and Skv must have the same length.The
[object Object]input supports the S length of 0 in a batch. In this case, the[object Object]input is not supported.The
[object Object]input does not support padding. That is,[object Object]cannot contain a row of all 1s.The S length of a batch in
[object Object]can be 0. If the S length is 0, the[object Object]input is not supported. If the actual S length is [2,2,0,2,2], the value of[object Object]is [2,4,4,6,8].Ascend 950PR/Ascend 950DT:
- seed and offset take effect only when keepProb is less than 1.0. Otherwise, they do not take effect.
- When keepProb is less than 1.0, if dropMaskOptional is not nullptr, the input dropMask is used. Otherwise, the dropMask generated by seed and offset is used.
In TND format, some batches at the end do not participate in computation. In this case, you can pass 0s to the end of actual_seq_q_len and actual_seq_kv_len. For example, if the actual S length is [2, 3, 4, 5, 6] and the last two batches are not required to participate in computation, the input actual_seq_q_len is [2, 3, 4, 0, 0]. In this case, if prefixOptional needs to be passed, the same number of 0s must be passed to the end of prefixOptional, for example, [1, 1, 1, 0, 0].
The following example is for reference only. For details, see .