aclnnPromptFlashAttentionV3
Applicable Products
| Product | Supported |
|---|---|
| × | |
| √ | |
| √ | |
| × | |
| √ | |
| × |
Function
Description: A FlashAttention operator for full inference scenarios, supporting sparse optimization, actualSeqLengthsKv optimization, int8 quantization, and the innerPrecise parameter (for selecting high-precision or high-performance mode).
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 attribute elements of the input sample, obtained through spatial transformation of the input sample and can be unified into a single 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 typically scaled by dividing by the square root of , and each row is normalized by softmax. After multiplying with , an matrix is obtained.
Function Prototype
The operator execution API is a two-phase API. You must first call the "aclnnPromptFlashAttentionV3GetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnPromptFlashAttentionV3" API to perform the computation.
aclnnStatus aclnnPromptFlashAttentionV3GetWorkspaceSize(
const aclTensor *query,
const aclTensor *key,
const aclTensor *value,
const aclTensor *pseShift,
const aclTensor *attenMask,
const aclIntArray *actualSeqLengths,
const aclIntArray *actualSeqLengthsKv,
const aclTensor *deqScale1,
const aclTensor *quantScale1,
const aclTensor *deqScale2,
const aclTensor *quantScale2,
const aclTensor *quantOffset2,
int64_t numHeads,
double scaleValue,
int64_t preTokens,
int64_t nextTokens,
char *inputLayout,
int64_t numKeyValueHeads,
int64_t sparseMode,
int64_t innerPrecise,
const aclTensor *attentionOut,
uint64_t *workspaceSize,
aclOpExecutor **executor)aclnnStatus aclnnPromptFlashAttentionV3(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
const aclrtStream stream)aclnnPromptFlashAttentionV3GetWorkspaceSize
Parameters
Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor query Input Input Q in the formula. Keep the data type consistent with key and value. FLOAT16, BFLOAT16, INT8 ND 3-4 × key Input Input K in the formula. Must be consistent with the data type of query and value. FLOAT16, BFLOAT16, INT8 ND 3-4 × value Input Input V in the formula. Must have the same data type as query and key. FLOAT16, BFLOAT16, INT8 ND 3-4 × pseShift Input. Positional encoding. - nullptr can be passed when this feature is not used.
- For comprehensive constraints, see Constraints.
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 - actualSeqLengthsKv Input Valid sequence lengths of key/value in different batches. - nullptr can be passed if no sequence length is specified.
- For comprehensive constraints, see Constraints.
INT64 ND 1 - deqScale1 Input Dequantization factor after BMM1. - Supports per-tensor.
- Pass nullptr when this feature is not used.
UINT64, FLOAT32 ND 1 - quantScale1 Input Quantization factor before BMM2. - Supports per-tensor.
- Pass nullptr when this feature is not used.
FLOAT32 ND 1 - deqScale2 Input Dequantization factor after BMM2. - Supports per-tensor.
- nullptr can be passed when this feature is not used.
UINT64, FLOAT32 ND 1 - quantScale2 Input Quantization factor of the output. - Supports per-tensor and per-channel.
- Pass nullptr when this function is not used.
FLOAT32, BFLOAT16 ND 1 - quantOffset2 Input Quantization offset of the output. - Supports per-tensor and per-channel.
- If this feature is not used, nullptr can be passed in.
FLOAT32, BFLOAT16 ND 1 - numHeads Input Number of heads in the query. - INT64 ND 1 - scaleValue Input The reciprocal of the square root of d in the formula. - The data type must satisfy the data type deduction rules with the data type of query.
- It is recommended to pass 1.0 if not explicitly specified by the user.
DOUBLE - 1 - preTokens Input Indicates the number of previous tokens that attention needs to compute correlation with. If not explicitly specified, it is recommended to pass 2147483647. INT64 - 1 - nextTokens Input Indicates the number of subsequent tokens with which attention needs to compute correlation. Pass 0 if not explicitly specified. INT64 - 1 - inputLayout Input Specifies the data layout format of the input query, key, and value. - When no specific format is required, "BSH" is recommended.
- For comprehensive constraints, see Constraints.
STRING - - - numKeyValueHeads Input Number of heads in key and value. - If not explicitly specified, passing 0 is recommended, indicating that the number of heads in key/value equals that in query.
- For comprehensive constraints, see Constraints.
INT64 - - - sparseMode Input Sparse mode. For comprehensive constraints, see Constraints. INT64 - 1 - innerPrecise Input Indicates the selection between high precision and high performance. For comprehensive constraints, see Constraints. INT8 - 1 - attentionOut Output Output in the formula. - FLOAT16, BFLOAT16, INT8 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 operator 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 the passed parameter is a mandatory input, output, or mandatory attribute and 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 internally calls an NPU runtime interface.
aclnnPromptFlashAttentionV3
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 from the first-phase API aclnnPromptFlashAttentionV3GetWorkspaceSize. 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:
- aclnnPromptFlashAttentionV3 defaults to a deterministic implementation.
When this API is used with PyTorch, ensure that the CANN-related packages match the PyTorch-related packages in version.
Handling of empty input parameters: The operator internally determines whether the query parameter is empty. If it is empty, the operator returns directly. If the query parameter is a non-empty tensor and the key and value parameters are empty tensors (i.e., S2 is 0), attentionOut is filled with all zeros. When attentionOut is an empty tensor, the AscendCLNN framework handles it. Other input parameters marked as "nullable" in the parameter description above are not processed when they are null pointers.
The data layout format of query, key, and value supports interpretation from multiple dimensions, where B (Batch) indicates the input sample batch size, S (Seq-Length) indicates the input sample sequence length, H (Head-Size) indicates the hidden layer size, N (Head-Num) indicates the number of heads, D (Head-Dim) indicates the minimum unit size of the hidden layer, and D = H/N is satisfied. T indicates the cumulative sum of the sequence lengths of all batch input samples.
The functional usage restrictions for query, key, and value inputs are as follows:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT :The B-axis supports values 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 supports values less than or equal to 256.
S supports values less than or equal to 20971520 (20M). In some long sequence scenarios, if the computation workload is too large, the PFA operator may time out (AI Core error with errorStr: timeout or trap error). In such scenarios, S splitting is recommended. Note: The computation workload here is affected by B, S, N, D, etc. Larger values result in greater computation workload. Typical long sequence scenarios that may cause timeouts (i.e., 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 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 TND scenarios:T must be less than or equal to 65536.
N must be 8/16/32/64/128, and Q_N, K_N, and V_N must be equal.
Q_D and K_D equal 192, and V_D equals 128 or 192.
Only the BFLOAT16 data type is supported.
The sparse mode only supports sparse=0 without a mask, or sparse=3 with a mask.
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, and 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-aligned, or when Q_S and KV_S have unequal lengths, the atten_mask configuration is not supported; the D axis must be less than or equal to 512.
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The supported data types are FLOAT16, BFLOAT16, and INT8.Atlas inference accelerator cards : Only FLOAT16 is supported as the data type.
The usage restrictions for the pseShift function are as follows:
The input shape type must be (B, N, Q_S, KV_S) or (1, N, Q_S, KV_S), where Q_S is S in the query shape and KV_S is S in the key and value shapes.
Q_S must be greater than or equal to the S length of query, and KV_S must be greater than or equal to the S length of key.
For scenarios where the KV_S of pseShift is not 32-byte aligned, padding to 32 bytes is recommended to improve performance. The padding values are not restricted. If this feature is not used, nullptr can be passed in.
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : FLOAT16 and BFLOAT16 data types are supported. When pseShift is of the FLOAT16 type, query must be of the FLOAT16 or INT8 type. When pseShift is of the BFLOAT16 type, query must be of the BFLOAT16 type. When query, key, and value are FLOAT16 and pseShift exists, the high-precision mode is used by default, and the corresponding constraints are inherited from those of the high-precision mode.Atlas inference accelerator cards : Only nullptr is supported.
Constraints on the attenMask function are as follows:
For scenarios where KV_S of attenMask is not 32-byte aligned, it is recommended to pad to 32-byte alignment to improve performance, with the padded portion filled with 1.
The recommended shape inputs 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.
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.
The usage restrictions for the actualSeqLengths and actualSeqLengthsKv inputs are as follows:
For actualSeqLengths, if the sequence length is not specified, nullptr can be passed, indicating that the valid sequence length is the same as the S length in the query shape. Note that the valid sequence length of each batch in this parameter must not exceed the sequence length of the corresponding batch in the query.
For actualSeqLengthsKv, if the sequence length is not specified, nullptr can be passed, indicating that the valid sequence length is the same as the S length in the key/value shape. Note that the valid sequence length of each batch in this parameter must not exceed the sequence length of the corresponding batch in the key/value.
The following rules apply to the input length of seqlen: When the input length is 1, all batches use the same seqlen. When the input length is greater than or equal to the number of batches, the first Batch values of seqlen are used. Other input lengths are not supported.
When the inputLayout of query is TND, this input parameter must be passed, and the number of elements in this input parameter is used as the Batch value. The value of each element in this input parameter represents the sum of the Sequence Lengths of the current batch and all previous batches. Therefore, the value of a subsequent element must be greater than or equal to the value of the preceding element, and negative values are not allowed.
Atlas A2 training products/Atlas A2 inference products : The data type supports INT64, and the TND format is supported.Atlas inference accelerator cards : Supports the INT64 data type.Ascend 950PR/Ascend 950DT : Supports the INT64 data type.
deqScale1 and deqScale2 inputs are subject to the following usage restrictions:
Atlas A2 training products/Atlas A2 inference products andAscend 950PR/Ascend 950DT : Support the UINT64 and FLOAT32 data types.Atlas inference accelerator cards : Only nullptr is supported.
quantScale1 input: The usage restrictions are as follows:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports FLOAT32.Atlas inference accelerator cards : Only nullptr is supported.
quantScale2 and quantOffset2 inputs: The usage restrictions are as follows:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports FLOAT32 and BFLOAT16.Atlas inference accelerator cards : Only nullptr is supported.
preTokens input, usage restrictions are as follows:
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.
nextTokens input, usage restrictions are as follows:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports INT64.Atlas inference accelerator cards : Only values 0 and 2147483647 are supported.
inputLayout input, with the following functional usage restrictions:
Currently supports BSH, BSND, BNSD, and BNSD_BSND (when the input is BNSD, the output format is BSND). If not explicitly specified, "BSH" is recommended.
Atlas A2 training products/Atlas A2 inference products : In addition to the above formats, TND is also supported (does not support pse, full quantization, or post-quantization).
numKeyValueHeads input, functional usage restrictions are as follows:
numHeads must be divisible by numKeyValueHeads, and in BSND, BNSD, and BNSD_BSND scenarios, it must be the same as the N-axis shape value of key/value in the shape, otherwise an error is reported.
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The data type supports INT64.Atlas inference accelerator cards : Only the value 0 is supported.
sparseMode input, functional usage restrictions are as follows:
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT :When sparseMode is 0, it indicates the defaultMask mode. If attenmask is not passed, no mask operation is performed, and preTokens and nextTokens are ignored (internally set to INT_MAX). If attenmask is passed, the complete attenmask matrix (S1 * S2) must be passed, indicating that the portion between preTokens and nextTokens needs to be computed.
When sparseMode is 1, it indicates the allMask mode, and the complete attenmask matrix (S1 * S2) must be passed.
When sparseMode is 2, it indicates the leftUpCausal mode mask, and the optimized attenmask matrix (2048*2048) must be passed.
When sparseMode is 3, it indicates the rightDownCausal mode mask, corresponding to the lower-triangular scenario with the right vertex as the partition. The optimized attenmask matrix (2048*2048) must be passed.
When sparseMode is 4, it indicates the band mode mask, which requires the optimized attenmask matrix (2048*2048).
When sparseMode is 5, 6, 7, or 8, they indicate prefix, global, dilated, and block_local respectively, none of which are currently supported. If the user does not specify a value, 0 is recommended.
Atlas A2 training products/Atlas A2 inference products : When inputLayout is TND, sparseMode only supports values 0 and 3.Atlas inference accelerator cards : Only value 0 is supported.
The innerPrecise input has the following usage restrictions:
There are four modes: 0, 1, 2, and 3. They occupy two bits in total. Bit 0 indicates the selection between high precision and high performance, and bit 1 indicates whether to perform row invalidity correction.
<table style="undefined;table-layout: fixed; width: 600px"><colgroup> <col style="width: 200px"> <col style="width: 200px"> <col style="width: 200px"> </colgroup><thead> <tr> <th>innerPrecise</th> <th>Mode</th> <th>Row Invalidity Correction</th> </tr></thead> <tbody> <tr> <td>0</td> <td>High precision</td> <td>×</td> </tr> <tr> <td>1</td> <td>High performance</td> <td>×</td> </tr> <tr> <td>2</td> <td>High precision</td> <td>√</td> </tr> <tr> <td>3</td> <td>High performance</td> <td>√</td> </tr> </tbody> </table>When Q_S > 1, sparse_mode is 0 or 1, and a user-defined mask is passed in, it is recommended to enable row invalidity.
BFLOAT16 and INT8 do not distinguish between high precision and high performance. Row invalidity correction takes effect for FLOAT16, BFLOAT16, and INT8. Currently, 0 and 1 are reserved configuration values. When the "mask portion involved in the computation" contains an entire row of all ones during calculation, precision loss may occur. In this case, you can try setting this parameter to 2 or 3 to enable the row invalidity function to improve precision, but this configuration will cause performance degradation. If the operator can determine that an invalid row scenario exists, it automatically enables invalid row computation, for example, when sparse_mode is 3 and Sq > Skv.
The usage restrictions for the attentionOut output are as follows:
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.
Atlas A2 training products/Atlas A2 inference products ,Ascend 950PR/Ascend 950DT : The supported data types are FLOAT16, BFLOAT16, and INT8.Atlas inference accelerator cards : Only FLOAT16 is supported.
Comprehensive constraints on the number of INT8 quantization-related parameters and the input/output data formats:
When the input is INT8 and the output is INT8: The parameters deqScale1, quantScale1, deqScale2, and quantScale2 must all be present. quantOffset2 is optional and defaults to 0 if not passed.
When the input is INT8 and the output is FLOAT16: The parameters deqScale1, quantScale1, and deqScale2 must all be present. If quantOffset2 or quantScale2 is present (i.e., not nullptr), an error is reported and returned.
When the input is FLOAT16 or BFLOAT16 and the output is INT8: The input parameter quantScale2 must be present, and quantOffset2 is optional (defaults to 0 if not passed). If deqScale1, quantScale1, or deqScale2 is present (i.e., not nullptr), an error is reported and the API returns.
The input parameters quantScale2 and quantOffset2 support both per-tensor and per-channel formats, and both FLOAT32 and BFLOAT16 data types. If quantOffset2 is passed, its type and shape information must be consistent with quantScale2. When the input is BFLOAT16, both FLOAT32 and BFLOAT16 are supported; otherwise, only FLOAT32 is supported. For the per-channel format, when the output layout is BSH, the product of all dimensions of quantScale2 must equal H; for other layouts, the product must equal N*D. (When the output layout is BSH, it is recommended to pass quantScale2 shape as [1,1,H] or [H]; when the output is BNSD, [1,N,1,D] or [N,D] is recommended; when the output is BSND, [1,1,N,D] or [N,D] is recommended.)
When the output is INT8 and quantScale2 and quantOffset2 are per-channel, left padding, Ring Attention, or scenarios where D is not 32-byte aligned are not supported.
When the output is INT8, the sparse mode with band and negative preTokens/nextTokens is not supported.
When the output is INT8, the input parameter quantOffset2 is passed with a non-null pointer and non-null tensor value, and sparseMode, preTokens, and nextTokens meet the following conditions, some rows of the matrix will not participate in the computation, leading to computation result errors. This scenario is intercepted. (Workaround: If you want this scenario not to be intercepted, perform post-quantization outside the PFA API instead of enabling it inside the PFA API.)
When sparseMode = 0, if attenMask is a non-null pointer and each batch satisfies actualSeqLengths - actualSeqLengthsKV - preTokens > 0 or nextTokens < 0, the interception condition is met.
When sparseMode = 1 or 2, the interception condition is never met.
When sparseMode = 3, if each batch satisfies actualSeqLengthsKV - actualSeqLengths < 0, the interception condition is met.
When sparseMode = 4, if preTokens < 0 or each batch satisfies nextTokens + actualSeqLengthsKV - actualSeqLengths < 0, the interception condition is met.
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 <cmath>
#include <cstring>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_prompt_flash_attention_v3.h"
#include "securec.h"
using namespace std;
namespace {
#define CHECK_RET(cond) ((cond) ? true :(false))
#define LOG_PRINT(message, ...) \
do { \
(void)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) {
auto ret = aclInit(nullptr);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclInit failed. ERROR: %d\n", ret);
return ret;
}
ret = aclrtSetDevice(deviceId);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret);
return ret;
}
ret = aclrtCreateStream(stream);
if (!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);
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret);
return ret;
}
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret);
return ret;
}
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];
}
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
shape.data(), shape.size(), *deviceAddr);
return 0;
}
struct TensorResources {
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;
aclIntArray* actualSeqLengths = nullptr;
};
int InitializeTensors(TensorResources& resources) {
std::vector<int64_t> queryShape = {1, 2, 1, 16};
std::vector<int64_t> keyShape = {1, 2, 2, 16};
std::vector<int64_t> valueShape = {1, 2, 2, 16};
std::vector<int64_t> attenShape = {1, 1, 1, 2};
std::vector<int64_t> outShape = {1, 2, 1, 16};
int64_t queryShapeSize = GetShapeSize(queryShape);
int64_t keyShapeSize = GetShapeSize(keyShape);
int64_t valueShapeSize = GetShapeSize(valueShape);
int64_t attenShapeSize = GetShapeSize(attenShape);
int64_t outShapeSize = GetShapeSize(outShape);
std::vector<float> queryHostData(queryShapeSize, 1);
std::vector<float> keyHostData(keyShapeSize, 1);
std::vector<float> valueHostData(valueShapeSize, 1);
std::vector<float> attenHostData(attenShapeSize, 1);
std::vector<float> outHostData(outShapeSize, 1);
int ret = CreateAclTensor(queryHostData, queryShape, &resources.queryDeviceAddr,
aclDataType::ACL_FLOAT16, &resources.queryTensor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
return ret;
}
ret = CreateAclTensor(keyHostData, keyShape, &resources.keyDeviceAddr,
aclDataType::ACL_FLOAT16, &resources.keyTensor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
return ret;
}
ret = CreateAclTensor(valueHostData, valueShape, &resources.valueDeviceAddr,
aclDataType::ACL_FLOAT16, &resources.valueTensor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
return ret;
}
ret = CreateAclTensor(attenHostData, attenShape, &resources.attenDeviceAddr,
aclDataType::ACL_BOOL, &resources.attenTensor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
return ret;
}
ret = CreateAclTensor(outHostData, outShape, &resources.outDeviceAddr,
aclDataType::ACL_FLOAT16, &resources.outTensor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
return ret;
}
std::vector<int64_t> actualSeqlenVector = {2};
resources.actualSeqLengths = aclCreateIntArray(actualSeqlenVector.data(),
actualSeqlenVector.size());
return ACL_SUCCESS;
}
int ExecutePromptFlashAttention(TensorResources& resources, aclrtStream stream,
void** workspaceAddr, uint64_t* workspaceSize) {
int64_t numHeads = 2;
int64_t numKeyValueHeads = numHeads;
double scaleValue = 1 / sqrt(2);
int64_t preTokens = 65535;
int64_t nextTokens = 65535;
int64_t sparseMode = 0;
int64_t innerPrecise = 1;
constexpr const char LAYER_OUT_STR[] = "BNSD";
constexpr size_t LAYER_OUT_LEN = sizeof(LAYER_OUT_STR);
char layerOut[LAYER_OUT_LEN];
memcpy(layerOut, LAYER_OUT_STR, LAYER_OUT_LEN);
aclOpExecutor* executor;
int ret = aclnnPromptFlashAttentionV3GetWorkspaceSize(
resources.queryTensor, resources.keyTensor, resources.valueTensor,
nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
numHeads, scaleValue, preTokens, nextTokens, layerOut, numKeyValueHeads,
sparseMode, innerPrecise, resources.outTensor, workspaceSize, &executor);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclnnPromptFlashAttentionV3GetWorkspaceSize failed. ERROR: %d\n", ret);
return ret;
}
if (*workspaceSize > 0ULL) {
ret = aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret);
return ret;
}
}
ret = aclnnPromptFlashAttentionV3(*workspaceAddr, *workspaceSize, executor, stream);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclnnPromptFlashAttentionV3 failed. ERROR: %d\n", ret);
return ret;
}
return ACL_SUCCESS;
}
int ProcessResults(TensorResources& resources, const std::vector<int64_t>& outShape) {
auto size = GetShapeSize(outShape);
std::vector<op::fp16_t> resultData(size, 0);
int ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
resources.outDeviceAddr, size * sizeof(resultData[0]),
ACL_MEMCPY_DEVICE_TO_HOST);
if (!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;
}
return ACL_SUCCESS;
}
void CleanupResources(TensorResources& resources, void* workspaceAddr,
aclrtStream stream, int32_t deviceId) {
if (resources.queryTensor) {
aclDestroyTensor(resources.queryTensor);
}
if (resources.keyTensor) {
aclDestroyTensor(resources.keyTensor);
}
if (resources.valueTensor) {
aclDestroyTensor(resources.valueTensor);
}
if (resources.attenTensor) {
aclDestroyTensor(resources.attenTensor);
}
if (resources.outTensor) {
aclDestroyTensor(resources.outTensor);
}
if (resources.actualSeqLengths) {
aclDestroyIntArray(resources.actualSeqLengths);
}
if (resources.queryDeviceAddr) {
aclrtFree(resources.queryDeviceAddr);
}
if (resources.keyDeviceAddr) {
aclrtFree(resources.keyDeviceAddr);
}
if (resources.valueDeviceAddr) {
aclrtFree(resources.valueDeviceAddr);
}
if (resources.attenDeviceAddr) {
aclrtFree(resources.attenDeviceAddr);
}
if (resources.outDeviceAddr) {
aclrtFree(resources.outDeviceAddr);
}
if (workspaceAddr) {
aclrtFree(workspaceAddr);
}
if (stream) {
aclrtDestroyStream(stream);
}
aclrtResetDevice(deviceId);
aclFinalize();
}
} // namespace
int main() {
int32_t deviceId = 0;
aclrtStream stream = nullptr;
TensorResources resources = {};
void* workspaceAddr = nullptr;
uint64_t workspaceSize = 0;
std::vector<int64_t> outShape = {1, 2, 1, 16};
int ret = ACL_SUCCESS;
// 1. Initialize device and stream
ret = Init(deviceId, &stream);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("Init acl failed. ERROR: %d\n", ret);
return ret;
}
// 2. Initialize tensors
ret = InitializeTensors(resources);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
CleanupResources(resources, workspaceAddr, stream, deviceId);
return ret;
}
// 3. Execute the operation
ret = ExecutePromptFlashAttention(resources, stream, &workspaceAddr, &workspaceSize);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
CleanupResources(resources, workspaceAddr, stream, deviceId);
return ret;
}
// 4. Synchronize stream
ret = aclrtSynchronizeStream(stream);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
CleanupResources(resources, workspaceAddr, stream, deviceId);
return ret;
}
// 5. Process results
ret = ProcessResults(resources, outShape);
if (!CHECK_RET(ret == ACL_SUCCESS)) {
CleanupResources(resources, workspaceAddr, stream, deviceId);
return ret;
}
// 6. Cleanup resources
CleanupResources(resources, workspaceAddr, stream, deviceId);
return 0;
}