API function: Updates the outputs of two FlashAttention operations based on their respective maximum and sum softmax values.
Formula:
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.
[object Object][object Object]
Parameters:
[object Object](aclTensor*, compute input): aclTensor on the device, prev_attn_out in the formula, output of the first FlashAttention operation. The data type can be FLOAT16, FLOAT, or BFLOAT16. The input shape must be the same as the[object Object]attribute. . The can be ND. When[object Object]is TND, D must be a multiple of 64.[object Object](aclTensor*, compute input): aclTensor on the device, prev_softmax_max in the formula, Softmax maximum result of the first FlashAttention operation. The data type can be FLOAT. The input shape is (B, N, S, 8) or (T, N, 8). The eight numbers of the last dimension must be identical and positive. . The can be ND. B indicates the batch size, N indicates the head number, S indicates the sequence length, and T indicates the time.[object Object](aclTensor*, compute input): aclTensor on the device, prev_softmax_sum in the formula, Softmax sum result of the first FlashAttention operation. The data type can be FLOAT. The input shape is the same as that of[object Object]. The eight numbers of the last dimension must be identical and positive. . The can be ND.[object Object](aclTensor*, compute input): aclTensor on the device, cur_attn_out in the formula, output of the second FlashAttention operation. The data type can be FLOAT16, FLOAT, or BFLOAT16. The data type and input shape must be the same as those of[object Object]. . The can be ND. When[object Object]is TND, D must be a multiple of 64.[object Object](aclTensor*, compute input): aclTensor on the device, cur_softmax_max in the formula, Softmax max result of the second FlashAttention operation. The data type can be FLOAT. The input shape is the same as that of[object Object]. The eight numbers of the last dimension must be identical and positive. . The can be ND.[object Object](aclTensor*, compute input): aclTensor on the device, cur_softmax_sum in the formula, Softmax sum result of the second FlashAttention operation. The data type can be FLOAT. The input shape is the same as that of[object Object]. The eight numbers of the last dimension must be identical and positive. . The can be ND.[object Object](aclTensor*, compute input): aclTensor on the device, cumulative sequence lengths starting from 0. The data type can be INT64. This parameter must be provided when[object Object]is TND. It is an integer-type aclTensor whose values increase from 0 to T.[object Object](char*, compute input): Host-side char* constant specifying the data layout of inputs related to attn_out.". Currently, TND and SBH are supported.[object Object](aclTensor*, compute output): aclTensor on the device, attn_out in the formula, updated output after both results are merged. The data type can be FLOAT16, FLOAT, or BFLOAT16. The data type and output shape must be the same as those of[object Object]. . The can be ND.[object Object](aclTensor*, compute output): aclTensor on the device, softmax_max in the formula, updated Softmax maximum after both results are merged. The data type can be FLOAT. The output shape must be the same as those of[object Object]. . The can be ND.[object Object](aclTensor*, compute output): aclTensor on the device, softmax_sum in the formula, updated Softmax sum after both results are merged. The data type can be FLOAT. The output shape must be the same as those of[object Object]. . The can be ND.[object Object](uint64_t*, output): size of the workspace to be allocated on the device.[object Object](aclOpExecutor**, output): operator executor, containing the operator computation process.
Returns:
[object Object]
Parameters:
[object Object](void*, input): address of the workspace to be allocated on the device.[object Object](uint64_t, input): size of the workspace to be allocated on the device, which is obtained by calling the first-phase API[object Object].[object Object](aclOpExecutor*, input): operator executor, containing the operator computation process.[object Object](aclrtStream, input): stream for executing the task.
Returns:
- Deterministic computing:
[object Object]defaults to a deterministic implementation.- When
[object Object]is TND, the last dimension of[object Object]must be a multiple of 64. - When
[object Object]is TND,[object Object]is required. - When
[object Object]is TND, N must be less than or equal to 256, and D must be less than or equal to 768.
The following example is for reference only. For details, see .
[object Object]