aclnnIncreFlashAttentionV4

📄 View Source Code

Applicable Products

Function

• Description: Compatible with the (aclnnIncreFlashAttentionV3) API, with the addition of the KV left-padding feature.

  For auto-regressive language models, as new tokens are generated, the inference input length keeps increasing. Incremental inference is implemented on top of the original full inference. The S axis of the query is fixed to 1. The key and value, after KV Cache, superimpose the previously inferred state information. The actual length of the S axis for each batch may differ, and the input data is fixed-length data after padding.

Compared with the FlashAttention operator in the full-attention scenario (PromptFlashAttention), the incremental inference process is not completely equivalent to normal full-attention inference, but the precision of incremental inference does not exhibit significant degradation.

NOTE

KV Cache is a common technique for optimizing the inference performance of large models. During sampling, the Transformer model uses the given prompt/context as the initial input for inference (which can be processed in parallel), and then generates additional tokens one by one to further refine the generated sequence (reflecting the autoregressive nature of the model). During the sampling process, the Transformer performs self-attention operations, which requires extracting key-value (KV) vectors for each item in the current sequence (whether prompt/context or generated tokens). These vectors are stored in a matrix, commonly referred to as the KV cache.

• Formula:

Self-attention leverages the relationships within the input sample itself to construct an attention model. The principle is as follows: assuming an input sample sequence $x$ of length $n$, where each element of $x$ is a $d$-dimensional vector, each $d$-dimensional vector can be regarded as a token embedding. This sequence is transformed through three weight matrices to obtain three matrices of dimension $n*d$.

The calculation formula of self-attention is generally defined as follows, where $Q$, $K$, and $V$ 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.

In this operator, the Score function uses the Softmax function, and the self-attention formula is:

Where the product of $Q$ and $K^T$ represents the attention of input $x$. To prevent this value from becoming too large, it is typically divided by the square root of $d$ for scaling, and each row is normalized by softmax. After multiplying with $V$, an $n*d$ matrix is obtained.

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnIncreFlashAttentionV4GetWorkspaceSize" API to obtain the workspace size required for computation and the executor that includes the operator computation process, and then call the "aclnnIncreFlashAttentionV4" API to execute the computation.

aclnnStatus aclnnIncreFlashAttentionV4GetWorkspaceSize(
    const aclTensor     *query,
    const aclTensorList *key,
    const aclTensorList *value,
    const aclTensor     *pseShift,
    const aclTensor     *attenMask,
    const aclIntArray   *actualSeqLengths,
    const aclTensor     *dequantScale1,
    const aclTensor     *quantScale1,
    const aclTensor     *dequantScale2,
    const aclTensor     *quantScale2,
    const aclTensor     *quantOffset2,
    const aclTensor     *antiquantScale,
    const aclTensor     *antiquantOffset,
    const aclTensor     *blocktable,
    const aclTensor     *kvPaddingSize,
    int64_t              numHeads,
    double               scaleValue,
    char                *inputLayout,
    int64_t              numKeyValueHeads,
    int64_t              blockSize,
    int64_t              innerPrecise,
    const aclTensor     *attentionOut,
    uint64_t            *workspaceSize,
    aclOpExecutor       **executor)

aclnnStatus aclnnIncreFlashAttentionV4(
    void             *workspace,
    uint64_t          workspaceSize,
    aclOpExecutor    *executor,
    const aclrtStream stream)

aclnnIncreFlashAttentionV4GetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  query	Input	Input Q in the formula.	The shapes of query and attentionOut must be fully consistent. Only supports the S axis of query equal to 1.	FLOAT16, BFLOAT16	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  key	Input	Input K in the formula.	The shapes of the corresponding tensors in key and value must be fully consistent.	FLOAT16, BFLOAT16, INT8	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  value	Input	Input V in the formula.	The shapes of the corresponding tensors in key and value must be fully consistent.	FLOAT16, BFLOAT16, INT8	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	×
  pseShift	Input	Positional encoding.	Supports null Tensor. The data type of pseShift must be consistent with that of query.	FLOAT16, BFLOAT16	ND	(B, N, 1, S) (1, N, 1, S)	×
  attenMask	Input	Attention mask matrix.	Supports null Tensor. When the data type of attenMask is INT8 or UINT8, the values in its tensor must be 0 or 1.	BOOL, INT8, UINT8	ND	(B, S) (B, 1, S) (B, 1, 1, S)	×
  actualSeqLengths	Input	Specifies the actual length of the S axis of key and value.	For comprehensive constraints, see Constraints.	INT64	ND	(B)	-
  dequantScale1	Input	Quantization factor for dequantization after BMM1.	Supports null Tensor. Supports per-tensor (scalar).	UINT64, FLOAT32	ND	(1)	×
  quantScale1	Input	Quantization factor for quantization before BMM2.	Supports null Tensor. Supports per-tensor (scalar).	FLOAT32, BFLOAT16	ND	(1)	×
  dequantScale2	Input	Quantization factor for dequantization after BMM2.	Supports null Tensor. Supports per-tensor (scalar).	UINT64, FLOAT32	ND	(1)	×
  quantScale2	Input	Quantization factor for output quantization.	Supports null Tensor. Supports per-tensor and per-channel.	FLOAT32, BFLOAT16	ND	(1) (D * N), where N is numKeyValueHeads.	×
  quantOffset2	Input	Quantization offset for output quantization.	Supports null Tensor. Supports per-tensor and per-channel.	FLOAT32, BFLOAT16	ND	(1) (D * N), N is numKeyValueHeads	×
  antiquantScale	Input	Quantization Factor.	Supports null Tensor. Supports per-tensor and per-channel (list), determined by shape. In BNSD scenario, the shape is (2, N, 1, D); in BSH scenario, the shape is (2, H); in BSND scenario, the shape is (2, N, D).	FLOAT16, BFLOAT16, FLOAT32	ND	(2, N, 1, D) (2, H) (2, N, D) (2)	×
  antiquantOffset	Input	Quantization factor for output quantization.	Supports null Tensor. Supports per-tensor and per-channel (list), determined by shape. In the BNSD scenario, the shape is (2, N, 1, D); in the BSH scenario, the shape is (2, H); in the BSND scenario, the shape is (2, N, D).	FLOAT16, BFLOAT16, FLOAT32	ND	(2, N, 1, D) (2, H) (2, N, D) (2)	×
  blocktable	Input	Block mapping table used for KV storage in page attention.	Supports null Tensor. For comprehensive constraints, see Constraints	INT32	ND	(B, maxBlockNumPerSeq)	×
  kvPaddingSize	Input	KV left-padding scenario.	Supports null Tensor.	INT64	-	-	×
  numHeads	Input	The number of heads in query.	numHeads is a multiple of numKeyValueHeads.	INT64	-	-	-
  scaleValue	Input	The reciprocal of the square root of d in the formula.	-	DOUBLE	-	-	-
  inputLayout	Input	Identifies the data layout format of the input query, key, and value.	Currently supports BSH, BNSD, and BSND. If the user does not specify it intentionally, it is recommended to pass "BSH".	STRING	-	-	-
  numKeyValueHeads	Input	Number of heads in key and value.	Used to support the GQA (Grouped-Query Attention) scenario. Passing 0 means the number is equal to the number of query heads. For comprehensive constraints, see Constraints.	INT64	-	-	-
  blockSize	Input	Specifies the maximum number of tokens stored in each block of KV in page attention.	If the user does not specify a value, 0 is recommended. For comprehensive constraints, see Constraints.	INT64	-	-	-
  innerPrecise	Input	Represents the high-precision/high-performance selection.	0 represents high precision, and 1 represents high performance. If the user does not specify a value, 1 is recommended.	INT64	-	-	-
  attentionOut	Output	Output in the formula.	-	FLOAT16, BFLOAT16, INT8	ND	(B, N, S, D) (B, S, N, D) (B, S, H)	-
  workspaceSize	Output	Returns the workspace size that the user needs to apply for on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator executor, which includes the operator computation process.	-	-	-	-	-

• 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	Returns 161001 if the input parameter is a required input, output, or required attribute and is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type and data format of query, key, value, actualSeqLengths, pseShift, attenMask, and attentionOut are not within the supported range.
  ACLNN_ERR_RUNTIME_ERROR	361001	The API memory call to the npu runtime interface is abnormal.

aclnnIncreFlashAttentionV4

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	The workspace memory address applied for on the device side.
  workspaceSize	Input	The workspace size applied for on the device side, obtained from the first-phase API aclnnIncreFlashAttentionV4GetWorkspaceSize.
  executor	Input	Operator executor that includes the calculation process of the operator.
  stream	Input	Specifies the stream for task execution.

• Return Value

  aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

• Deterministic computation:

  • aclnnIncreFlashAttentionV4 defaults to a deterministic implementation.

• Atlas A2 training products/Atlas A2 inference products, Ascend 950PR/Ascend 950DT:

  • Supports B-axis less than or equal to 65536, N-axis less than or equal to 256, and D-axis less than or equal to 512.

  • The query data type supports FLOAT16 and BFLOAT16, and the attentionOut, key, and value data types support FLOAT16, INT8, and BFLOAT16.

  • The dequantScale1 and dequantScale2 data types support UINT64 and FLOAT32.

  • The quantScale1, quantScale2, and quantOffset2 data types support FLOAT32.

  • The numKeyValueHeads data type supports INT64.

• Atlas inference products:

  • Supports B-axis less than or equal to 256, N-axis less than or equal to 256, and D-axis less than or equal to 512.

  • Supports key and value S-axis less than or equal to 65536.

  • query, key, value, and attentionOut Only Supports FLOAT16 data type.

  • dequantScale1, quantScale1, dequantScale2, quantScale2, and quantOffset2 Only Supports nullptr.

  • numKeyValueHeads Only Supports the value 0.

• In non-contiguous scenarios, the number of tensors in the tensorlist of the key and value parameters is equal to B of query (due to tensorlist limitations, B must be less than or equal to 256 in non-contiguous scenarios). The shape must be fully consistent except for S, and batch can only be 1.

• 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 dimension of the hidden layer, and D = H/N is satisfied.

• The N in the query parameter is equal to the numHeads value, the N of key and value is equal to the numKeyValueHeads value, and numHeads is a multiple of numKeyValueHeads.

• The scenario where query, key, and value inputs are all INT8 is not supported.

• In the scenario where query, key, and value inputs are FLOAT16 and the output is INT8: the input parameter quantScale2 is required, quantOffset2 is optional, and the dequantScale1, quantScale1, and dequantScale2 parameters cannot be passed (i.e., they must be nullptr).

• Constraints on antiquantScale and antiquantOffset parameters:

  • Per-channel mode: The shape of both parameters can be (2, N, 1, D), (2, N, D), or (2, H), where N is numKeyValueHeads. The data type of the parameters is the same as the query data type.

  • Per-tensor mode: The shape of both parameters is (2), and the data type is the same as the query data type.

• 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 be equal to H; for other layouts, the product must be equal to N*D. (When the output layout is BSH, it is recommended to pass [1,1,H] or [H] for the quantScale2 shape; 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.)

• Page attention scenario:

  • The prerequisite for enabling page attention is that blockTable exists and is valid, and key and value are arranged in a contiguous memory according to the indices in blockTable. key and value dtype supports FLOAT16/BFLOAT16/INT8. In this scenario, the inputLayout parameter of key and value is invalid.

  • blockSize is a user-defined parameter. Its value affects the performance of page attention. When page attention is enabled, blockSize must be a non-zero value, and the maximum blockSize does not exceed 512. When the key and value input type is FLOAT16/BFLOAT16, it must be 16-aligned. When the key and value input type is INT8, it must be 32-aligned. 128 is recommended. Typically, page attention can improve throughput but causes performance degradation.

  • In the page attention scenario, when the inputLayout of query is BNSD, the KV cache layout supports two formats: (blocknum, blocksize, H) and (blocknum, KV_N, blocksize, D). When the inputLayout of query is BSH or BSND, the KV cache layout supports only one format: (blocknum, blocksize, H). blocknum must not be less than the sum of the number of blocks for each batch calculated based on actualSeqLengthsKv and blockSize. In addition, the shapes of key and value must be consistent.

  • In the page attention scenario, the performance of the KV cache layout (blocknum, KV_N, blocksize, D) is generally better than that of the KV cache layout (blocknum, blocksize, H). It is recommended to prioritize the (blocknum, KV_N, blocksize, D) format.

  • In the page attention enabled scenario, when the input KV cache layout format is (blocknum, blocksize, H) and numKvHeads * headDim exceeds 64k, it will be intercepted and an error will be reported due to hardware instruction constraints. This can be resolved by enabling GQA (reducing numKvHeads) or adjusting the KV cache layout format to (blocknum, numKvHeads, blocksize, D).

  • In the page attention scenario, the input actualSeqLengths must be passed.

  • In the page attention scenario, blockTable must be two-dimensional, with the first dimension length equal to B and the second dimension length not less than maxBlockNumPerSeq (maxBlockNumPerSeq is the number of blocks corresponding to the maximum actualSeqLengthsKv in each batch).

  • In the page attention enabled scenario, the input S must be greater than or equal to maxBlockNumPerSeq * blockSize in the following scenarios.

  • When Attention mask is enabled, for example, the mask shape is (B, 1, 1, S).

  • When pseShift is enabled, for example, the pseShift shape is (B, N, 1, S).

  • Enable kvPaddingSize, for example, kvPaddingSize shape is (1), and S must be greater than or equal to kvPaddingSize plus the sum of the maximum actualSeqLengthsKv in each Batch.

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 "securec.h"
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_incre_flash_attention_v4.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;
    aclTensorList* tensorKeyList = nullptr;
    aclTensorList* tensorValueList = 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<int8_t> 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_INT8, &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;
    }

    int kvTensorNum = 1;
    aclTensor* tensorsOfKey[] = {resources.keyTensor};
    resources.tensorKeyList = aclCreateTensorList(tensorsOfKey, kvTensorNum);

    aclTensor* tensorsOfValue[] = {resources.valueTensor};
    resources.tensorValueList = aclCreateTensorList(tensorsOfValue, kvTensorNum);

    std::vector<int64_t> actualSeqlenVector = {2};
    resources.actualSeqLengths = aclCreateIntArray(actualSeqlenVector.data(),
                                                  actualSeqlenVector.size());

    return ACL_SUCCESS;
}

int ExecuteIncreFlashAttention(TensorResources& resources, aclrtStream stream,
                              void** workspaceAddr, uint64_t* workspaceSize) {
    int64_t numHeads = 2;
    int64_t numKeyValueHeads = numHeads;
    int64_t blockSize = 1;
    int64_t innerPrecise = 1;
    double scaleValue = 1 / sqrt(2);
    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 = aclnnIncreFlashAttentionV4GetWorkspaceSize(
        resources.queryTensor, resources.tensorKeyList, resources.tensorValueList,
        nullptr, resources.attenTensor, resources.actualSeqLengths, nullptr,
        nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
        numHeads, scaleValue, layerOut, numKeyValueHeads, blockSize, innerPrecise,
        resources.outTensor, workspaceSize, &executor);

    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("aclnnIncreFlashAttentionV4GetWorkspaceSize 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 = aclnnIncreFlashAttentionV4(*workspaceAddr, *workspaceSize, executor, stream);
    if (!CHECK_RET(ret == ACL_SUCCESS)) {
        LOG_PRINT("aclnnIncreFlashAttentionV4 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 = ExecuteIncreFlashAttention(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;
}