aclnnPromptFlashAttention
Note: This API will be deprecated in later versions. Use the latest API aclnnPromptFlashAttentionV3 instead.
Applicable Products
| Product | Supported |
|---|---|
| × | |
| × | |
| √ | |
| × | |
| √ | |
| × |
Function
Description: FlashAttention operator for full inference scenarios.
Formula:
Self-attention leverages the relationships within the input sample itself to construct an attention model. The principle is as follows: given an input sample sequence of length , where each element of is a -dimensional vector, each -dimensional vector can be regarded as a token embedding. This sequence is transformed through three weight matrices to obtain three matrices of dimension .
The calculation formula of self-attention is generally defined as follows, where , , and are important property elements of the input sample, obtained through spatial transformation of the input sample and can be unified into a common feature space. "Attention" in the formula and operator name is an abbreviation of "self-attention".
In this operator, the Score function uses the Softmax function, and the self-attention calculation formula is:
Where: the product of and represents the attention of the input . To prevent this value from becoming too large, it is usually scaled by dividing by the square root of , and each row is normalized by softmax. After multiplying with , an n*d matrix is obtained.
Function Prototype
The operator execution API is a two-phase API. You must first call the "aclnnPromptFlashAttentionGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnPromptFlashAttention" API to perform the computation.
aclnnStatus aclnnPromptFlashAttentionGetWorkspaceSize(
const aclTensor *query,
const aclTensor *key,
const aclTensor *value,
const aclTensor *pseShift,
const aclTensor *attenMask,
const aclIntArray *actualSeqLengths,
int64_t numHeads,
double scaleValue,
int64_t preTokens,
int64_t nextTokens,
char *inputLayout,
int64_t numKeyValueHeads,
const aclTensor *attentionOut,
uint64_t *workspaceSize,
aclOpExecutor **executor)aclnnStatus aclnnPromptFlashAttention(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
const aclrtStream stream)aclnnPromptFlashAttentionGetWorkspaceSize
Parameters
Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-Contiguous Tensor query Input Input Q in the formula. Must have the same data type as key and value. FLOAT16, BFLOAT16 ND 3-4 × key Input. Input K in the formula. Keep the data type consistent with query and value. FLOAT16, BFLOAT16 ND 3-4 × value Input Input V in the formula. Must have the same data type as query and key. FLOAT16, BFLOAT16 ND 3-4 × pseShift Input Positional encoding Reserved parameter, not used yet. Currently, this parameter is forcibly set to nullptr. FLOAT16, BFLOAT16 ND 4 × attenMask Input Mask matrix. - Pass nullptr if this feature is not used.
- For comprehensive constraints, see Constraints.
BOOL, INT8, UINT8 ND 2-4 × actualSeqLengths Input Effective sequence length of the query in different batches. - If the sequence length is not specified, nullptr can be passed in.
- For comprehensive constraints, see Constraints.
INT64 ND 1 - numHeads Input. Number of heads in query. Constraint: In BNSD/NSD scenarios, this value must be the same as the N-axis shape value of query in shape. INT64 ND 1 - scaleValue Input Reciprocal of the square root of d in the formula. - The data type must meet the data type deduction rules with the data type of query.
- If not specified by the user, 1.0 is recommended.
DOUBLE - 1 - preTokens Input Indicates the number of previous tokens that attention needs to compute association with. - If not specified, 2147483647 is recommended.
- Negative values are supported.
INT64 - 1 - nextTokens Input Indicates the number of subsequent tokens that attention needs to compute correlation with. - If not explicitly specified, 0 is recommended.
- Negative numbers are supported.
INT64 - 1 - inputLayout Input Indicates the data layout format of the input query, key, and value. - If not explicitly specified, "BSH" is recommended.
- For comprehensive constraints, see Constraints.
STRING - - - numKeyValueHeads Input Number of heads in key and value. - If not specified by the user, it is recommended to pass 0, indicating that the number of heads in key/value is the same as that in query.
- For comprehensive constraints, see Constraints.
INT64 - - - attentionOut Output Output in the formula. - FLOAT16, BFLOAT16 ND 3-4 - workspaceSize Output Returns the workspace size that the user needs to apply for on the Device side. - - - 1 - executor Output Returns the op executor, which contains the operator computation flow. - - - 1 - Return Value
aclnnStatus: return code. For details, see aclnn Return Code.
The first-phase API performs input parameter validation and returns an error in the following scenarios:
Return Value Error Code Description ACLNN_ERR_PARAM_NULLPTR 161001 If a mandatory input, output, or mandatory attribute is a null pointer, 161001 is returned. ACLNN_ERR_PARAM_INVALID 161002 The data type or data format of query, key, value, pseShift, attenMask, or attentionOut is not within the supported range. ACLNN_ERR_RUNTIME_ERROR 361001 An exception occurs when the API memory calls the npu runtime interface.
aclnnPromptFlashAttention
Parameters
Parameter Input/Output Description workspace Input Address of the workspace memory applied for on the Device side. workspaceSize Input Size of the workspace applied for on the Device side, obtained through the first-phase API aclnnPromptFlashAttentionGetWorkspaceSize. executor Input Operator executor, which contains the operator computation flow. stream Input Specifies the stream that executes the task. Return Value
aclnnStatus: return code. For details, see aclnn Return Code.
Constraints
Deterministic computation:
- aclnnPromptFlashAttention defaults to a deterministic implementation.
When this API is used with PyTorch, ensure that the versions of CANN-related packages match those of PyTorch-related packages.
Handling of null input parameters: The operator internally determines whether the query parameter is null. If it is null, the operator returns directly. If the query parameter is a non-null tensor but the key and value parameters are null tensors (that is, S2 is 0), attentionOut is filled with all zeros. When attentionOut is a null tensor, the AscendCLNN framework handles it. For other input parameters marked as "nullptr allowed" in the parameter description above, no processing is performed when they are null pointers.
The data layout format of query, key, and value supports interpretation from multiple dimensions, where B (Batch) indicates the batch size of input samples, S (Seq-Length) indicates the sequence length of input samples, H (Head-Size) indicates the hidden layer size, N (Head-Num) indicates the number of heads, and D (Head-Dim) indicates the minimum unit dimension of the hidden layer, satisfying D = H/N.
The functional usage restrictions for query, key, and value inputs are as follows:
Input shape constraints:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT :The B-axis must be less than or equal to 65536 (64k). When the input type includes INT8 and the D-axis is not 32-byte aligned, or when the input type is FLOAT16 or BFLOAT16 and the D-axis is not 16-byte aligned, the B-axis only supports up to 128.
The N-axis must be less than or equal to 256.
S supports values less than or equal to 20971520 (20M). In some long sequence scenarios, excessive computation may cause the PFA operator execution to time out (aicore error, with errorStr: timeout or trap error). In such scenarios, it is recommended to perform S splitting. Note: The computation amount here is affected by B, S, N, D, etc. Larger values result in greater computation. Typical long sequence scenarios that may cause timeout (i.e., scenarios where the product of B, S, N, and D is large) include but are not limited to:
B Q_N Q_S D KV_N KV_S 1 20 2097152 256 1 2097152 1 2 20971520 256 2 20971520 20 1 2097152 256 1 2097152 1 10 2097152 512 1 2097152 The D axis must be less than or equal to 512. When inputLayout is BSH or BSND, N*D must be less than 65535.
Atlas A2 training products/Atlas A2 inference products : Comprehensive constraints on query, key, and value inputs in the TND scenario:T must be less than or equal to 65536.
N equals 8/16/32/64/128, and Q_N, K_N, and V_N must be equal.
Q_D and K_D equal to 192, and V_D equal to 128 or 192;
Only the BFLOAT16 data type is supported;
The sparse mode only supports sparse=0 with no mask passed, or sparse=3 with a mask passed;
When sparse=3, each batch requires actualSeqLengths < actualSeqLengthsKv.
Atlas inference accelerator cards :When inputLayout is BSH, the B axis must be less than or equal to 300; in other cases, the B axis must be less than or equal to 128.
The N axis must be less than or equal to 256.
The S axis must be less than or equal to 65535 (64k). When Q_S or KV_S is not 128-byte aligned, or when Q_S and KV_S are not equal in length, configuring atten_mask is not supported.
The D axis must be less than or equal to 512.
Input data type constraints:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : FLOAT16 and BFLOAT16 data types are supported.Atlas inference accelerator cards : Only the FLOAT16 data type is supported.
The pseShift function usage constraints are as follows:
Input data type constraints:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : FLOAT16 and BFLOAT16 data types are supported.Atlas inference accelerator cards : Only nullptr is supported.
The usage restrictions of attenMask are as follows:
Input shape restrictions: Typically, the recommended shapes are Q_S, KV_S; B, Q_S, KV_S; 1, Q_S, KV_S; B, 1, Q_S, KV_S; 1, 1, Q_S, KV_S, where Q_S is S in the shape of query, and KV_S is S in the shape of key and value.
Input data type restrictions:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The supported data types are BOOL, INT8, and UINT8.Atlas inference accelerator cards : Only BOOL is supported.
Other constraints: For scenarios where KV_S of attenMask is not 32-byte aligned, it is recommended to pad to 32-byte alignment to improve performance, and fill the excess part with 1.
The functional usage constraints of the actualSeqLengths input are as follows:
Input value range constraint: The valid sequence length of each batch in this input parameter must not exceed the sequence length of the corresponding batch in query.
Input data type constraints:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports INT64.Atlas inference accelerator cards : The data type supports INT64.
The preTokens input has the following functional restrictions:
Input data type restrictions:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports INT64.Atlas inference accelerator cards : Only the value 2147483647 is supported.
The following describes the functional restrictions on the nextTokens input:
Input data type restrictions:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The INT64 data type is supported.Atlas inference accelerator cards : Only the values 0 and 2147483647 are supported.
The inputLayout parameter has the following usage restrictions: Input data type restrictions: Currently, BSH, BSND, BNSD, and BNSD_BSND are supported (when the input is in BNSD format, the output format is BSND). If you do not specify a format, BSH is recommended. The numKeyValueHeads parameter has the following usage restrictions: Input attribute restriction: numHeads must be divisible by numKeyValueHeads. In BSND, BNSD, and BNSD_BSND scenarios, numKeyValueHeads must be equal to the N-axis shape value of key/value in the shape. Otherwise, an error is reported.
Input data type restrictions:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : INT64 is supported.Atlas inference accelerator cards : Only the value 0 is supported.
attentionOut output: The usage restrictions are as follows:
Shape restrictions: When inputLayout is BNSD_BSND, the shape of the input query is BNSD and the output shape is BSND. In other cases, the shape of this input parameter must be consistent with the shape of the input query.
Data type constraints:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : FLOAT16 and BFLOAT16 are supported.Atlas inference accelerator cards : Only FLOAT16 is supported.
Example
The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.
#include <iostream>
#include <vector>
#include <math.h>
#include <cstring>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_prompt_flash_attention.h"
using namespace std;
#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;
}
int Init(int32_t deviceId, aclrtStream* stream) {
// Boilerplate for resource initialization.
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 = 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 apply for device-side memory.
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);
// Copy the host-side data to the device-side memory by calling aclrtMemcpy.
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);
// Calculate the strides of a 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 the aclCreateTensor API 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. (Boilerplate) Initialize the device/stream. For details, see the AscendCL API list.
// Set deviceId based on the actual device.
int32_t deviceId = 0;
aclrtStream stream;
auto ret = Init(deviceId, &stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
// 2. Construct the input and output. Customize the construction based on the API interface.
int32_t batchSize = 1;
int32_t numHeads = 2;
int32_t sequenceLengthQ = 1;
int32_t headDims = 16;
int32_t keyNumHeads = 2;
int32_t sequenceLengthKV = 16;
std::vector<int64_t> queryShape = {batchSize, numHeads, sequenceLengthQ, headDims}; // BNSD
std::vector<int64_t> keyShape = {batchSize, keyNumHeads, sequenceLengthKV, headDims}; // BNSD
std::vector<int64_t> valueShape = {batchSize, keyNumHeads, sequenceLengthKV, headDims}; // BNSD
std::vector<int64_t> attenShape = {batchSize, 1, 1, sequenceLengthKV}; // B11S
std::vector<int64_t> outShape = {batchSize, numHeads, sequenceLengthQ, headDims}; // BNSD
void *queryDeviceAddr = nullptr;
void *keyDeviceAddr = nullptr;
void *valueDeviceAddr = nullptr;
void *attenDeviceAddr = nullptr;
void *outDeviceAddr = nullptr;
aclTensor *queryTensor = nullptr;
aclTensor *keyTensor = nullptr;
aclTensor *valueTensor = nullptr;
aclTensor *attenTensor = nullptr;
aclTensor *outTensor = nullptr;
std::vector<float> queryHostData(batchSize * numHeads * sequenceLengthQ * headDims, 1.0f);
std::vector<float> keyHostData(batchSize * keyNumHeads * sequenceLengthKV * headDims, 1.0f);
std::vector<float> valueHostData(batchSize * keyNumHeads * sequenceLengthKV * headDims, 1.0f);
std::vector<int8_t> attenHostData(batchSize * sequenceLengthKV, 0);
std::vector<float> outHostData(batchSize * numHeads * sequenceLengthQ * headDims, 1.0f);
// Create the query aclTensor.
ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the key aclTensor.
ret = CreateAclTensor(keyHostData, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &keyTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the value aclTensor.
ret = CreateAclTensor(valueHostData, valueShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &valueTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the atten aclTensor.
ret = CreateAclTensor(attenHostData, attenShape, &attenDeviceAddr, aclDataType::ACL_BOOL, &attenTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
// Create the out aclTensor.
ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &outTensor);
CHECK_RET(ret == ACL_SUCCESS, return ret);
std::vector<int64_t> actualSeqlenVector = {sequenceLengthKV};
auto actualSeqLengths = aclCreateIntArray(actualSeqlenVector.data(), actualSeqlenVector.size());
int64_t numKeyValueHeads = numHeads;
double scaleValue = 1 / sqrt(headDims); // 1/sqrt(d)
int64_t preTokens = 65535;
int64_t nextTokens = 65535;
string sLayerOut = "BNSD";
char layerOut[sLayerOut.length()+1];
strcpy(layerOut, sLayerOut.c_str());
// 3. Call the CANN operator library API.
uint64_t workspaceSize = 0;
aclOpExecutor* executor;
// Call the first-phase API.
ret = aclnnPromptFlashAttentionGetWorkspaceSize(queryTensor, keyTensor, valueTensor, nullptr, nullptr, nullptr, numHeads, scaleValue,
preTokens, nextTokens, layerOut, numKeyValueHeads, outTensor, &workspaceSize, &executor);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnPromptFlashAttentionGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
// Apply for device memory based on the workspaceSize computed by the first-phase API.
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.
ret = aclnnPromptFlashAttention(workspaceAddr, workspaceSize, executor, stream);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnPromptFlashAttention failed. ERROR: %d\n", ret); return ret);
// 4. (Boilerplate) Synchronously wait for task execution to finish.
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 memory to the host memory. Modify this based on the specific API definition.
auto size = GetShapeSize(outShape);
std::vector<op::fp16_t> resultData(size, 0);
ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
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 ret);
for (int64_t i = 0; i < size; i++) {
std::cout << "index: " << i << ": " << static_cast<float>(resultData[i]) << std::endl;
}
// 6. Release resources.
aclDestroyTensor(queryTensor);
aclDestroyTensor(keyTensor);
aclDestroyTensor(valueTensor);
aclDestroyTensor(attenTensor);
aclDestroyTensor(outTensor);
aclDestroyIntArray(actualSeqLengths);
aclrtFree(queryDeviceAddr);
aclrtFree(keyDeviceAddr);
aclrtFree(valueDeviceAddr);
aclrtFree(attenDeviceAddr);
aclrtFree(outDeviceAddr);
if (workspaceSize > 0) {
aclrtFree(workspaceAddr);
}
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}