aclnnFlashAttentionUnpaddingScoreGrad
Supported Products
| Product | Supported |
|---|---|
| √ | |
| √ | |
| × | |
| × | |
| × | |
| × |
Description
API function: Computes the attention backpropagation output in training scenarios, which is the backpropagation of
aclnnFlashAttentionVarLenScore.Formula:
The forward propagation formula for attention is as follows:
For convenience, the formula can be represented using variables and :
Then the backward propagation formula for attention is as follows:
Prototype
Each operator has two-phase API calls. First, aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, aclnnFlashAttentionUnpaddingScoreGrad is called to perform computation.
aclnnStatus aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize(
const aclTensor* query,
const aclTensor* keyIn,
const aclTensor* value,
const aclTensor* dy,
const aclTensor* pseShiftOptional,
const aclTensor* dropMaskOptional,
const aclTensor* paddingMaskOptional,
const aclTensor* attenMaskOptional,
const aclTensor* softmaxMaxOptional,
const aclTensor* softmaxSumOptional,
const aclTensor* softmaxInOptional,
const aclTensor* attentionInOptional,
const aclIntArray* prefixOptional,
const aclIntArray* actualSeqQLenOptional,
const aclIntArray* actualSeqKvLenOptional,
double scaleValue,
double keepProb,
int64_t preTokens,
int64_t nextTokens,
int64_t headNum,
char* inputLayout,
int64_t innerPrecise,
int64_t sparseMode,
const aclTensor* dqOut,
const aclTensor* dkOut,
const aclTensor* dvOut,
const aclTensor* dpseOut,
uint64_t* workspaceSize,
aclOpExecutor** executor)aclnnStatus aclnnFlashAttentionUnpaddingScoreGrad(
void* workspace,
uint64_t workspaceSize,
aclOpExecutor* executor,
const aclrtStream stream)aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize
Parameters
Name Input/Output Description Usage Notes Data Type Data Format Dimension (Shape) Non-contiguous Tensor query Input Q in the formula. The data type must be the same as that of keyIn or value. FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ keyIn Input K in the formula. The data type must be the same as that of query or value. FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ value Input V in the formula. The data type must be the same as that of query or keyIn. FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ dy Input dY in the formula. - FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ pseShiftOptional Optional input pse in the formula. The data type must match query. Use this parameter with pseType. FLOAT16, BFLOAT16, FLOAT32 ND [B,N,1024,Skv], [1,N,1024,Skv] √ dropMaskOptional Input Dropout mask. - UINT8 ND 0, 1 √ paddingMaskOptional Input Reserved parameter. A null pointer must be passed to this parameter when the API is called. - - - - attenMaskOptional Input atten_mask in the formula. A value of 1 indicates that the position does not participate in the calculation, while a value of 0 indicates that it does. BOOL, UINT8 ND [B,N,Sq,Skv], [B,1,Sq,Skv], [1,1,Sq,Skv], [Sq,Skv] √ softmaxMaxOptional Input Intermediate output of the softmax forward propagation. - FLOAT ND [N,T,8] √ softmaxSumOptional Input Intermediate output of the softmax forward propagation. - FLOAT ND [N,T,8] √ softmaxInOptional Input Intermediate output of the softmax forward propagation. Reserved. - - - - attentionInOptional Input Forward attention output. The data type and shape must be the same as those of query. FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ prefixOptional Input N of each batch in the prefix sparse computation scenario. - INT64 ND 0, 1 - actualSeqQLenOptional Input Actual query sequence length. - INT64 ND 1 - actualSeqKvLenOptional Input Actual key/value sequence length. - INT64 ND 1 - dqOut Output dQ in the formula, indicating the gradient of query. - FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ dkOut Output dK in the formula, indicating the gradient of key. - FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ dvOut Output dV in the formula, indicating the gradient of value. - FLOAT16, BFLOAT16, FLOAT32 ND [TND] √ dpseOut Output d(pse) gradient. Reserved. - - - - scaleValue Input Scale factor. - DOUBLE - - - keepProb Input Ratio of 1s in dropMaskOptional. - DOUBLE - - - preTokens Input Left boundary of the sliding window for sparse computation. - INT64 - - - nextTokens Input Right boundary of the sliding window for sparse computation. - INT64 - - - headNum Input Number of heads on a single rank, that is, the length of the N axis of query. - INT64 - - - inputLayout Input Data layout of input Q/K/V. TND is supported. String - - - innerPrecise Input Internal calculation precision control. Reserved. INT64 - - - sparseMode Input Sparse mode. The value ranges from 0 to 8, excluding 5. INT64 - - - workspaceSize Output Size of the workspace required to be allocated on the device. - - - - - executor Output Operator executor, containing the operator computation process. - - - - - Returns
aclnnStatus: status code. For details, see aclnn Return Codes.
The first-phase API implements input parameter verification. The following errors may be thrown.
Return Error Code Description ACLNN_ERR_PARAM_NULLPTR 161001 The required input, output, or attribute is passed as a null pointer. ACLNN_ERR_PARAM_INVALID 161002 The data type of query, keyIn, value, dy, pseShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOptional, softmaxSumOptional, softmaxInOptional, attentionInOptional, dqOut, dkOut, or dvOut is not supported. The data format of query, keyIn, value, dy, pseShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOptional, softmaxSumOptional, softmaxInOptional, attentionInOptional, dqOut, dkOut, or dvOut is not supported.
aclnnFlashAttentionUnpaddingScoreGrad
Parameters
Name Input/Output Description workspace Input Address of the workspace to be allocated on the device. workspaceSize Input Size of the workspace to be allocated on the device, which is obtained by calling the first-phase API aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize. executor Input Operator executor, containing the operator computation process. stream Input Stream for executing the task. Returns
aclnnStatus status code. For details, see aclnn Return Codes.
Constraints
- Deterministic computing:
aclnnFlashAttentionUnpaddingScoreGraddefaults to a non-deterministic implementation. You can callaclrtCtxSetSysParamOptto enable deterministic computing.
- When this API is used together with PyTorch, ensure that the CANN package versions match the PyTorch package versions.
- B (batch size) of the input
query,key,value, anddymust be the same.inputLayoutmust be consistent. - D (Head-Dim) of the input
query,key, andvaluemust satisfy (qD == kD && kD >= vD). - The input data types of
query,key,value, andpseShiftOptionalmust be the same. - The shapes of the input
keyandvaluemust be the same. If the D values ofquery,key, andvalueare the same, the shapes ofqueryanddymust be the same. - N of the input
querycan be different from N of thekeyorvalue, but they must be proportional. That is, Nq/Nkv must be a non-zero integer and the value of Nq ranges from 1 to 256. - The following uses the
inputLayoutTND as an example to describe the constraints on the shape:- T: The value ranges from 1 to 1M.
- N: The value ranges from 1 to 256.
- D: The value ranges from 1 to 768.
- KeepProb: The value range is (0, 1].
- The data format of
query,key, andvaluecan only be TND. T indicates the data closely arranged on the B and S axes (SeqLenQ and SeqLenKV of each batch). B (Batch) indicates the batch size of the input sample, and 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. D (Head-Dim) indicates the minimum unit size of the hidden layer (D = H/N). realShiftOptional: If Sq is greater than 1024, Sq and Skv of each batch are of equal length, and it is a lower triangular mask scenario withsparseModebeing 0, 2, or 3, ALiBi positional encoding compression can be enabled. In this case, only the last 1024 rows of the original PSE need to be input for memory optimization, that is,alibi_compress = ori_pse[:, :, -1024:, :]. Specifically:- If the parameters of each batch are different, the shape is BNHSkv (H=1024).
- When each batch is the same, the shape is 1NHSkv (H=1024).
- If this parameter is not used, a null pointer can be passed.
- The constraints for
sparseModeare as follows:- If the shape values of all
attenMaskOptionalare the same and less than 2048, you are advised to use the default mode to reduce memory usage. - When the value is set to 1, 2, 3, or 5, the user-configured
preTokensandnextTokensdo not take effect. - When the value is set to 0 or 4, ensure that the ranges of
attenMaskOptional,preTokens, andnextTokensare consistent. - If no specific value is required, you are advised to set it to 0.
- For details about the sparse modes, see Sparse Mode Description.
- When the value is set to 7,
realShiftOptionalis not supported. - When the value is set to 8,
realShiftOptionalis supported when the q and kv of each sequence have the same length. PSE generation is performed globally. Outer splitting in the q direction is supported. q and kv of each sequence must have the same length before outer splitting, andactualSeqQLenOptionalis passed after outer splitting.
- If the shape values of all
- In some scenarios, if the computation load is too large, the operator execution may time out (AI Core error, errorStr: timeout or trap error). 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.
- The
prefixOptionalsparse computing supports only compression scenarios (sparseMode = 6). When Sq > Skv, the value range of N ofprefixis [0, Skv]. When Sq ≤ Skv, the value range of N ofprefixis [Skv – Sq, Skv].[0]-actualSeqKvLenOptional[0]+qStartIdxOptional-kvStartIdxOptional== 0 (experimental feature) - The
actualSeqQLenOptionalinput supports the S length of 0 in a batch. In this case, thepseShiftOptionalinput is not supported. The length ofactualSeqQLenOptionalranges from 1 to 2K. WhenprefixOptionalis present, its maximum length is 1K. - Constraints on the
softmaxMaxandsoftmaxSumparameters: The input format is fixed at [B, N, S, 8], except TND format, which is [T, N, 8]. Note: T = B x S. - The value of
headNummust be the same as the value of N inquery.
Example
The following example is for reference only. For details, see Compilation and Running Sample.
#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_flash_attention_score_grad.h"
#define CHECK_RET(cond, return_expr) \
do { \
if (!(cond)) { \
return_expr; \
} \
} while (0)
#define LOG_PRINT(message, ...) \
do { \
printf(message, ##__VA_ARGS__); \
} while (0)
int64_t GetShapeSize(const std::vector<int64_t>& shape) {
int64_t shapeSize = 1;
for (auto i : shape) {
shapeSize *= i;
}
return shapeSize;
}
void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
auto size = GetShapeSize(shape);
std::vector<float> resultData(size, 0);
auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
*deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
for (int64_t i = 0; i < size; i++) {
LOG_PRINT("mean result[%ld] is: %f\n", i, resultData[i]);
}
}
int Init(int32_t deviceId, aclrtContext* context, aclrtStream* stream) {
// (Fixed writing) Initialize resources.
auto ret = aclInit(nullptr);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
ret = aclrtSetDevice(deviceId);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
ret = aclrtCreateContext(context, deviceId);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
ret = aclrtSetCurrentContext(*context);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
ret = aclrtCreateStream(stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
return 0;
}
template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
aclDataType dataType, aclTensor** tensor) {
auto size = GetShapeSize(shape) * sizeof(T);
// Call aclrtMalloc to allocate memory on the device.
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
// Call aclrtMemcpy to copy the data on the host to the memory on the device.
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
// Compute the strides of the contiguous tensor.
std::vector<int64_t> strides(shape.size(), 1);
for (int64_t i = shape.size() - 2; i >= 0; i--) {
strides[i] = shape[i + 1] * strides[i + 1];
}
// Call aclCreateTensor to create an aclTensor.
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
shape.data(), shape.size(), *deviceAddr);
return 0;
}
int main() {
// 1. (Fixed writing) Initialize the device and stream. For details, see the ACL API manual.
// Set the device ID in use.
int32_t deviceId = 0;
aclrtContext context;
aclrtStream stream;
auto ret = Init(deviceId, &context, &stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
// 2. Construct the inputs and outputs based on the API definition.
std::vector<int64_t> qShape = {256, 1, 128};
std::vector<int64_t> kShape = {256, 1, 128};
std::vector<int64_t> vShape = {256, 1, 128};
std::vector<int64_t> dxShape = {256, 1, 128};
std::vector<int64_t> attenmaskShape = {256, 256};
std::vector<int64_t> softmaxMaxShape = {256, 1, 8};
std::vector<int64_t> softmaxSumShape = {256, 1, 8};
std::vector<int64_t> attentionInShape = {256, 1, 128};
std::vector<int64_t> dqShape = {256, 1, 128};
std::vector<int64_t> dkShape = {256, 1, 128};
std::vector<int64_t> dvShape = {256, 1, 128};
void* qDeviceAddr = nullptr;
void* kDeviceAddr = nullptr;
void* vDeviceAddr = nullptr;
void* dxDeviceAddr = nullptr;
void* attenmaskDeviceAddr = nullptr;
void* softmaxMaxDeviceAddr = nullptr;
void* softmaxSumDeviceAddr = nullptr;
void* attentionInDeviceAddr = nullptr;
void* dqDeviceAddr = nullptr;
void* dkDeviceAddr = nullptr;
void* dvDeviceAddr = nullptr;
aclTensor* q = nullptr;
aclTensor* k = nullptr;
aclTensor* v = nullptr;
aclTensor* dx = nullptr;
aclTensor* pse = nullptr;
aclTensor* dropMask = nullptr;
aclTensor* padding = nullptr;
aclTensor* attenmask = nullptr;
aclTensor* softmaxMax = nullptr;
aclTensor* softmaxSum = nullptr;
aclTensor* softmaxIn = nullptr;
aclTensor* attentionIn = nullptr;
aclTensor* dq = nullptr;
aclTensor* dk = nullptr;
aclTensor* dv = nullptr;
aclTensor* dpse = nullptr;
std::vector<float> qHostData(32768, 1);
std::vector<float> kHostData(32768, 1);
std::vector<float> vHostData(32768, 1);
std::vector<float> dxHostData(32768, 1);
std::vector<uint8_t> attenmaskHostData(65536, 0);
std::vector<float> softmaxMaxHostData(2048, 3.0);
std::vector<float> softmaxSumHostData(2048, 3.0);
std::vector<float> attentionInHostData(32768, 1);
std::vector<float> dqHostData(32768, 0);
std::vector<float> dkHostData(32768, 0);
std::vector<float> dvHostData(32768, 0);
ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_FLOAT16, &dx);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_FLOAT16, &attentionIn);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_FLOAT16, &dk);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_FLOAT16, &dv);
CHECK_RET(ret == ACL_SUCCESS, return ret);
std::vector<int64_t> prefixOp = {0};
aclIntArray* prefix = aclCreateIntArray(prefixOp.data(), 1);
std::vector<int64_t> acSeqQLenOp = {256};
std::vector<int64_t> acSeqKvLenOp = {256};
aclIntArray* acSeqQLen = aclCreateIntArray(acSeqQLenOp.data(), acSeqQLenOp.size());
aclIntArray* acSeqKvLen = aclCreateIntArray(acSeqKvLenOp.data(), acSeqKvLenOp.size());
double scaleValue = 0.088388;
double keepProb = 1;
int64_t preTokens = 65536;
int64_t nextTokens = 65536;
int64_t headNum = 1;
int64_t innerPrecise = 0;
int64_t sparseMode = 0;
char layOut[5] = {'T', 'N', 'D', 0};
// 3. Call the CANN operator library API, which needs to be replaced with the actual API.
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// Call the first-phase API of aclnnFlashAttentionUnpaddingScoreGrad.
ret = aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize(q, k, v, dx, pse, dropMask, padding,
attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, prefix, acSeqQLen, acSeqKvLen,
scaleValue, keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise, sparseMode,
dq, dk, dv, dpse, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionUnpaddingScoreGradGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
// Allocate device memory based on the computed workspaceSize.
void* workspaceAddr = nullptr;
if (workspaceSize > 0) {
ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
}
// Call the second-phase API of aclnnFlashAttentionUnpaddingScoreGrad.
ret = aclnnFlashAttentionUnpaddingScoreGrad(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionUnpaddingScoreGrad failed. ERROR: %d\n", ret); return ret);
// 4. (Fixed writing) Wait until the task execution is complete.
ret = aclrtSynchronizeStream(stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
// 5. Obtain the output value and copy the result from the device to the host. Modify the code based on the API definition.
PrintOutResult(dqShape, &dqDeviceAddr);
PrintOutResult(dkShape, &dkDeviceAddr);
PrintOutResult(dvShape, &dvDeviceAddr);
// 6. Release aclTensor and aclScalar. Modify the configuration based on the API definition.
aclDestroyTensor(q);
aclDestroyTensor(k);
aclDestroyTensor(v);
aclDestroyTensor(dx);
aclDestroyTensor(attenmask);
aclDestroyTensor(softmaxMax);
aclDestroyTensor(softmaxSum);
aclDestroyTensor(attentionIn);
aclDestroyTensor(dq);
aclDestroyTensor(dk);
aclDestroyTensor(dv);
// 7. Release device resources.
aclrtFree(qDeviceAddr);
aclrtFree(kDeviceAddr);
aclrtFree(vDeviceAddr);
aclrtFree(dxDeviceAddr);
aclrtFree(attenmaskDeviceAddr);
aclrtFree(softmaxMaxDeviceAddr);
aclrtFree(softmaxSumDeviceAddr);
aclrtFree(attentionInDeviceAddr);
aclrtFree(dqDeviceAddr);
aclrtFree(dkDeviceAddr);
aclrtFree(dvDeviceAddr);
if (workspaceSize > 0) {
aclrtFree(workspaceAddr);
}
aclrtDestroyStream(stream);
aclrtDestroyContext(context);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}