aclnnBlockSparseAttention

Applicable Products

Function

• Description: BlockSparseAttention performs sparse attention computation, supporting flexible block-level sparse patterns. It uses BlockSparseMask to specify the KV blocks selected for each Q block, enabling efficient sparse attention computation.

• Formula: Sparse block size: $blockShapeX \times blockShapeY$, selectIdx specifies the sparse pattern

  $$attentionOut = Softmax(scale \cdot query \cdot key_{sparse}^T + atten\_mask) \cdot value_{sparse}$$

  The data arrangement format of the query, key, and value inputs for BlockSparseAttention supports interpretation from multiple dimension arrangements, which can be passed in through qInputLayout and kvInputLayout.

  • B: indicates the batch size of the input samples (Batch).

  • T: the length of B and S tightly packed along the combined axis (Total tokens).

  • S: indicates the sequence length of the input samples (Seq-Length).

  • H: indicates the size of the hidden layer (Head-Size).

  • N: indicates the number of heads (Head-Num).

  • D: indicates the minimum unit size of the hidden layer, which must satisfy D = H/N (Head-Dim).

  Currently supported layouts:

  • qInputLayout: "TND" "BNSD"

  • kvInputLayout: "TND" "BNSD"

Function Prototype

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

aclnnStatus aclnnBlockSparseAttentionGetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *key,
  const aclTensor   *value,
  const aclTensor   *blockSparseMaskOptional,
  const aclTensor   *attenMaskOptional,
  const aclIntArray *blockShapeOptional,
  const aclIntArray *actualSeqLengthsOptional,
  const aclIntArray *actualSeqLengthsKvOptional,
  const aclTensor   *blockTableOptional,
  char              *qInputLayout,
  char              *kvInputLayout,
  int64_t            numKeyValueHeads,
  int64_t            maskType,
  double             scaleValue,
  int64_t            innerPrecise,
  int64_t            blockSize,
  int64_t            preTokens,
  int64_t            nextTokens,
  int64_t            softmaxLseFlag,
  aclTensor         *attentionOut,
  aclTensor         *softmaxLseOptional,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnBlockSparseAttentionGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-Contiguous Tensor
  query	Input	The query in the formula.	Supported shapes: TND: [totalQTokens, headNum, headDim]. BNSD: [batch, headNum, maxQSeqLength, headDim].	FLOAT16, BFLOAT16	ND	3/4	×
  key	Input	The key in the formula.	Supported shapes: TND: [totalKTokens, numKeyValueHeads, headDim]. BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim].	FLOAT16, BFLOAT16	ND	3/4	×
  value	Input	value in the formula.	Supported shapes: TND: [totalVTokens, numKeyValueHeads, headDim]. BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim].	FLOAT16, BFLOAT16	ND	3/4	×
  blockSparseMaskOptional	Input	Indicates the actual sparse pattern.	Optional input (mandatory in the current version). Shape: [batch, headNum, ceilDiv(maxQSeqLength, blockShapeX), ceilDiv(maxKvSeqLength, blockShapeY)]. Indicates which blocks participate in computation (1) and which do not (0) after block partitioning. If nullptr is passed in, block sparse computation is disabled, meaning attention scores between all tokens are computed.	INT8	ND	4	×
  attenMaskOptional	Input	atten_mask in the formula.	atten_mask is superimposed with the sparse pattern to take effect. Currently not supported; nullptr must be passed in.	INT8	ND	2	×
  blockShapeOptional	Input	aclIntArray on the host side. An array of sparse block shapes.	Used together with blockSparseMaskOptional: When blockSparseMaskOptional is configured: if this input is configured, the operator obtains the sparse block size from it; if not configured, the operator defaults the sparse block size to [128, 128]. When blockSparseMaskOptional is not configured: the operator ignores this input regardless of its configuration. When this input is configured: it must contain two elements [blockShapeX, blockShapeY]. blockShapeX: block size in the Q direction. The value must be greater than 0. blockShapeY: block size in the KV direction. The value must be greater than 0 and a multiple of 128.	INT64	-	1	-
  actualSeqLengthsOptional	Input	An aclIntArray on the host side that describes the query sequence length corresponding to each batch.	Optional input, used for variable-length sequence scenarios: When qInputLayout is "TND": This input must be configured. When qInputLayout is "BNSD": If this input is configured, within the operator processing is performed based on the actual sequence length specified by this input; if not configured (nullptr passed in), within the operator processing is performed based on S in the shape of query.	INT64	-	1	-
  actualSeqLengthsKvOptional	Input	aclIntArray on the host side, describing the key/value sequence length corresponding to each batch.	Optional input, used for variable-length sequence scenarios: When kvInputLayout is "TND": this input must be configured. When kvInputLayout is "BNSD": if this input is configured, the operator processes data according to the actual sequence length specified by this input; if this input is not configured (nullptr is passed in), the operator processes data according to S in the shape of key/value.	INT64	-	1	-
  blockTableOptional	Input	Block table used for PagedAttention.	Currently not supported. Must be passed in as nullptr.	INT32	ND	2	×
  qInputLayout	Input	A string on the host side, indicating the data arrangement format of the input query.	Currently only "TND" and "BNSD" are supported. qInputLayout must be consistent with kvInputLayout.	String	-	-	-
  kvInputLayout	Input	String on the host side, indicating the data arrangement format of the input key and value.	Currently only supports "TND" and "BNSD". qInputLayout and kvInputLayout must be consistent.	String	-	-	-
  numKeyValueHeads	Input	An int64_t value on the host side, indicating the number of heads for key/value.	-	INT64	-	-	-
  maskType	Input	int64_t on the host side, which indicates the mask type in attention computation.	Currently only supports passing 0. 0: indicates the no-mask scenario.	INT64	-	-	-
  scaleValue	Input	A double on the host side, representing the scaling factor in the formula.	Generally set to D^-0.5.	DOUBLE	-	-	-
  innerPrecise	Input	int64_t on the host side, indicating the precision level for Softmax computation.	Controls the data type used for operations in the online softmax phase and the rescale phase. Currently only supports passing 0 or 1. 0: Indicates that both online softmax and rescale use the fp32 data type, suitable for scenarios that prioritize computation precision. 1: Only supported when the input query, key, and value are all of the fp16 data type. Indicates that both online softmax and rescale use the fp16 data type, which delivers better performance but lower precision and may cause numerical overflow during computation. Users must determine whether to use this option based on the value range.	INT64	-	-	-
  blockSize	Input	int64_t on the host side, indicating the block size of PagedAttention.	Used for the PagedAttention scenario. PagedAttention is currently not supported, so only passing 0 is supported.	INT64	-	-	-
  preTokens	Input	An int64_t value on the host side. Indicates the number of tokens the sliding window needs to include forward in a sliding window attention scenario.	Used for sliding window attention scenarios. Sliding window attention is currently not supported; only 2147483647 can be passed in.	INT64	-	-	-
  nextTokens	Input	int64_t on the host side. In the sliding window attention scenario, indicates how many tokens the sliding window needs to include backward.	Used in the sliding window attention scenario. Sliding window attention is currently not supported; only supports passing in 2147483647.	INT64	-	-	-
  softmaxLseFlag	Input	A flag of type int64_t on the host side that indicates whether to enable softmaxLse output.	Currently only supports passing 0 or 1. 0: Indicates that softmaxLse is not output. 1: Indicates that softmaxLse is output. Performance loss may occur compared to not outputting softmaxLse.	INT64	-	-	-
  attentionOut	Output	attentionOut in the formula.	The data type and shape are consistent with query.	FLOAT16, BFLOAT16	ND	3/4	√
  softmaxLseOptional	Output	Intermediate log-sum-exp result of the Softmax computation.	Supported shapes vary with the shape of query: query is "TND": [totalQTokens, headNum, 1]. query is "BNSD": [batch, headNum, maxQSeqLength, 1].	FLOAT	ND	3/4	√
  workspaceSize	Output	Returns the workspace size to be applied on the device side.	-	-	-	-	-
  executor	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

• 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 Code	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The input query, key, or value is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type of query, key, or value is not within the supported range.
  		qInputLayout or kvInputLayout is invalid.
  		blockShape is invalid (the number of elements is less than 2 or a value is less than or equal to 0).
  		innerPrecise is invalid (must be 0, 1, or 4).

aclnnBlockSparseAttention

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Workspace memory address applied on the device side.
  workspaceSize	Input	Workspace size applied on the device side, obtained from the first-phase API aclnnBlockSparseAttentionGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the AscendCL stream that executes the task.

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnBlockSparseAttention defaults to a deterministic implementation.

• When this API is used with PyTorch, ensure that the versions of CANN-related packages and PyTorch-related packages match.

• qInputLayout currently only supports "TND" and "BNSD".

• kvInputLayout currently only supports "TND" and "BNSD".

• Currently, the InputLayout of query, key, and value must be consistent.

• The data types of the input query, key, and value must be consistent, and support FLOAT16 and BFLOAT16.

• The D axis of query, key, and value currently only supports being configured as 64 or 128.

• If blockShapeOptional is passed in, it must contain at least two elements [blockShapeX, blockShapeY], the values must be greater than 0, and blockShapeY must be a multiple of 128.

• blockSparseMaskOptional currently must be passed in, and its shape must be [batch, headNum, ceilDiv(maxQS, blockShapeX), ceilDiv(maxKVS, blockShapeY)].

• attentionMaskOptional currently only supports passing nullptr.

• actualSeqLengthsOptional is mandatory when qInputLayout is "TND"; actualSeqLengthsKvOptional is mandatory when kvInputLayout is "TND".

• actualSeqLengthsOptional and actualSeqLengthsKvOptional currently must be configured together or not configured together. Configuring only one of them will be intercepted by the operator.

• blockTableOptional currently only supports passing nullptr, which indicates that the PagedAttention feature is not enabled.

• innerPrecise must be 0 (float32 softmax) or 1 (fp16 softmax). When the query input is BFLOAT16, it can only be configured as 0.

• qSeqlen and kvSeqlen do not need to be divisible by blockShape. Unaligned scenarios are supported, and the actual number of blocks is calculated by rounding up.

• If the headNum of the input query is N1 and the headNum of the input key and value is N2, then N1 >= N2 && N1 % N2 == 0.

• maskType currently only supports input 0, which indicates that no mask is applied.

• blockSize currently only supports input 0, which indicates that paged cache is not supported.

• preTokens and nextTokens currently only support passing in 2147483647, which indicates that all tokens before and after the current token participate in the attention computation, i.e., sliding window attention is not 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 <cstring>
#include <cmath>
#include <cstdint>
#include "acl/acl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_block_sparse_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, AscendCL 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) {
    // Check whether the shape is valid.
    if (shape.empty()) {
        LOG_PRINT("CreateAclTensor: ERROR - shape is empty\n");
        return -1;
    }
    for (size_t i = 0; i < shape.size(); ++i) {
        if (shape[i] <= 0) {
            LOG_PRINT("CreateAclTensor: ERROR - shape[%zu]=%ld is invalid\n", i, shape[i]);
            return -1;
        }
    }

    auto size = GetShapeSize(shape) * sizeof(T);

    // Check whether the hostData size matches.
    if (hostData.size() != static_cast<size_t>(GetShapeSize(shape))) {
        LOG_PRINT("CreateAclTensor: ERROR - hostData size mismatch: %zu vs %ld\n",
                  hostData.size(), GetShapeSize(shape));
        return -1;
    }

    // Call aclrtMalloc to apply for device-side memory.
    *deviceAddr = nullptr;
    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 data from the host side 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);

    // Compute the strides of a contiguous tensor.
    std::vector<int64_t> strides(shape.size(), 1);
    if (shape.size() > 1) {
        for (int64_t i = static_cast<int64_t>(shape.size()) - 2; i >= 0; i--) {
            strides[i] = shape[i + 1] * strides[i + 1];
        }
    }

    // Create an aclTensor by calling the aclCreateTensor API.
    *tensor = nullptr;
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    CHECK_RET(*tensor != nullptr, LOG_PRINT("aclCreateTensor failed - returned nullptr\n");
            return -1);
    return 0;
}

void FreeResource(aclTensor *query, aclTensor *key, aclTensor *value, aclTensor *blockSparseMask,
                  aclTensor *attentionOut, aclIntArray *actualSeqLengths, aclIntArray *actualSeqLengthsKv,
                  aclIntArray *blockShape, void *queryDeviceAddr, void *keyDeviceAddr, void *valueDeviceAddr,
                  void *blockSparseMaskDeviceAddr, void *attentionOutDeviceAddr, void *actualSeqLengthsDeviceAddr,
                  void *actualSeqLengthsKvDeviceAddr, void *workspaceAddr, int32_t deviceId, aclrtStream *stream)
{
    // Release resources.
    if (query) {
        aclDestroyTensor(query);
    }
    if (key) {
        aclDestroyTensor(key);
    }
    if (value != nullptr) {
        aclDestroyTensor(value);
    }
    if (blockSparseMask) {
        aclDestroyTensor(blockSparseMask);
    }
    if (attentionOut) {
        aclDestroyTensor(attentionOut);
    }
    if (actualSeqLengths) {
        aclDestroyIntArray(actualSeqLengths);
    }
    if (actualSeqLengthsKv) {
        aclDestroyIntArray(actualSeqLengthsKv);
    }
    if (blockShape) {
        aclDestroyIntArray(blockShape);
    }

    if (queryDeviceAddr) {
        aclrtFree(queryDeviceAddr);
    }
    if (keyDeviceAddr) {
        aclrtFree(keyDeviceAddr);
    }
    if (valueDeviceAddr) {
        aclrtFree(valueDeviceAddr);
    }
    if (blockSparseMaskDeviceAddr) {
        aclrtFree(blockSparseMaskDeviceAddr);
    }
    if (attentionOutDeviceAddr) {
        aclrtFree(attentionOutDeviceAddr);
    }
    if (actualSeqLengthsDeviceAddr) {
        aclrtFree(actualSeqLengthsDeviceAddr);
    }
    if (actualSeqLengthsKvDeviceAddr) {
        aclrtFree(actualSeqLengthsKvDeviceAddr);
    }
    if (workspaceAddr) {
        aclrtFree(workspaceAddr);
    }

    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
}

int main() {
    // 1. (Boilerplate) Initialize the device/stream.
    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. Set parameters.
    int32_t batch = 1;
    int32_t qSeqlen = 128;
    int32_t kvSeqlen = 128;
    int32_t numHeads = 1;
    int32_t numKvHeads = 1;
    int32_t headDim = 128;
    int32_t blockShapeX = 128;
    int32_t blockShapeY = 128;

    // Calculate TND format dimensions.
    int64_t totalQTokens = batch * qSeqlen;
    int64_t totalKvTokens = batch * kvSeqlen;
    int32_t qBlockNum = (qSeqlen + blockShapeX - 1) / blockShapeX;  // Number of Q blocks in the X dimension.
    int32_t kvBlockNum = (kvSeqlen + blockShapeY - 1) / blockShapeY;  // Number of KV blocks in the Y dimension.
    // totalQBlocks = qBlockNum * numHeads (each Q block corresponds to one head).
    int32_t totalQBlocks = qBlockNum * numHeads;
    int32_t maxKvBlockNum = kvBlockNum;

    aclTensor *queryTensor = nullptr;
    aclTensor *keyTensor = nullptr;
    aclTensor *valueTensor = nullptr;
    aclTensor *blockSparseMaskTensor = nullptr;
    aclTensor *attentionOutTensor = nullptr;
    aclIntArray *actualSeqLengths = nullptr;
    aclIntArray *actualSeqLengthsKv = nullptr;
    aclIntArray *blockShape = nullptr;

    void *queryDeviceAddr = nullptr;
    void *keyDeviceAddr = nullptr;
    void *valueDeviceAddr = nullptr;
    void *blockSparseMaskDeviceAddr = nullptr;
    void *attentionOutDeviceAddr = nullptr;
    void *actualSeqLengthsDeviceAddr = nullptr;
    void *actualSeqLengthsKvDeviceAddr = nullptr;
    void* workspaceAddr = nullptr;

    // 3. Create the Query tensor (TND format: [totalQTokens, numHeads, headDim]).
    std::vector<int64_t> queryShape = {totalQTokens, numHeads, headDim};
    std::vector<op::fp16_t> queryHostData(totalQTokens * numHeads * headDim, 1.0f);
    ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to create query tensor\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 4. Create the Key/Value tensor (TND format: [totalKvTokens, numKvHeads, headDim]).
    std::vector<int64_t> kvShape = {totalKvTokens, numKvHeads, headDim};
    std::vector<op::fp16_t> keyHostData(totalKvTokens * numKvHeads * headDim, 1.0f);
    std::vector<op::fp16_t> valueHostData(totalKvTokens * numKvHeads * headDim, 1.0f);
    ret = CreateAclTensor(keyHostData, kvShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &keyTensor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to create key tensor\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);
    ret = CreateAclTensor(valueHostData, kvShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &valueTensor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to create value tensor\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 5. Create the blockSparseMask tensor ([batch, numHeads, qBlockNum, kvBlockNum]).
    std::vector<int8_t> blockSparseMaskHostData(totalQBlocks * numHeads, 0);
    blockSparseMaskHostData[0] = static_cast<int8_t>(1);
    std::vector<int64_t> blockSparseMaskShape = {batch, numHeads, qBlockNum, kvBlockNum};
    ret = CreateAclTensor(blockSparseMaskHostData, blockSparseMaskShape, &blockSparseMaskDeviceAddr, aclDataType::ACL_INT8, &blockSparseMaskTensor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to create block sparse mask tensor\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 6. Create the output tensor.
    std::vector<int64_t> attentionOutShape = {totalQTokens, numHeads, headDim};
    int64_t attentionOutElementCount = totalQTokens * numHeads * headDim;
    std::vector<op::fp16_t> attentionOutHostData(attentionOutElementCount, 0.0f);
    ret = CreateAclTensor(attentionOutHostData, attentionOutShape, &attentionOutDeviceAddr, aclDataType::ACL_FLOAT16, &attentionOutTensor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to create attentionOut tensor\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 7. Create the blockShape array.
    std::vector<int64_t> blockShapeData = {blockShapeX, blockShapeY};
    blockShape = aclCreateIntArray(blockShapeData.data(), blockShapeData.size());
    CHECK_RET(blockShape != nullptr, LOG_PRINT("Failed to create blockShape array\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return -1);

    // 8. Create actualSeqLengths and actualSeqLengthsKv (required parameters).
    std::vector<int64_t> actualSeqLengthsHost(batch, static_cast<int64_t>(qSeqlen));
    std::vector<int64_t> actualSeqLengthsKvHost(batch, static_cast<int64_t>(kvSeqlen));

    size_t seqLengthsSize = batch * sizeof(int64_t);

    ret = aclrtMalloc(&actualSeqLengthsDeviceAddr, seqLengthsSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to allocate actualSeqLengths memory\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);
    ret = aclrtMalloc(&actualSeqLengthsKvDeviceAddr, seqLengthsSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to allocate actualSeqLengthsKv memory\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    ret = aclrtMemcpy(actualSeqLengthsDeviceAddr, seqLengthsSize, actualSeqLengthsHost.data(), seqLengthsSize,
                      ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to copy actualSeqLengths to device\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);
    ret = aclrtMemcpy(actualSeqLengthsKvDeviceAddr, seqLengthsSize, actualSeqLengthsKvHost.data(), seqLengthsSize,
                      ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Failed to copy actualSeqLengthsKv to device\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // aclCreateIntArray expects a data pointer on the host side, not on the device side.
    actualSeqLengths = aclCreateIntArray(actualSeqLengthsHost.data(), batch);
    actualSeqLengthsKv = aclCreateIntArray(actualSeqLengthsKvHost.data(), batch);
    CHECK_RET(actualSeqLengths != nullptr && actualSeqLengthsKv != nullptr,
              LOG_PRINT("Failed to create actualSeqLengths arrays\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return -1);

    // 9. Prepare string parameters (ensure the buffer size is sufficient, including the null terminator).
    const char* qLayoutStr = "TND";
    const char* kvLayoutStr = "TND";
    char qLayoutBuffer[16] = {0};
    char kvLayoutBuffer[16] = {0};
    strncpy(qLayoutBuffer, qLayoutStr, sizeof(qLayoutBuffer) - 1);
    strncpy(kvLayoutBuffer, kvLayoutStr, sizeof(kvLayoutBuffer) - 1);

    // 10. Calculate scaleValue.
    float scaleValue = 1.0f / std::sqrt(static_cast<float>(headDim));

    // 11. Call the first-phase API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;

    ret = aclnnBlockSparseAttentionGetWorkspaceSize(
        queryTensor,           // query
        keyTensor,             // key
        valueTensor,           // value
        blockSparseMaskTensor, // blockSparseMask
        nullptr,               // attenMaskOptional
        blockShape,            // blockShape
        actualSeqLengths,      // actualSeqLengthsOptional
        actualSeqLengthsKv,    // actualSeqLengthsKvOptional
        nullptr,               // blockTableOptional
        qLayoutBuffer,         // qInputLayout
        kvLayoutBuffer,        // kvInputLayout
        numKvHeads,            // numKeyValueHeads
        0,                     // maskType
        scaleValue,            // scaleValue
        0,                     // innerPrecise (1=fp16 softmax)
        0,                     // blockSize
        2147483647,            // preTokens
        2147483647,            // nextTokens
        0,                     // softmaxLseFlag
        attentionOutTensor,          // attentionOut
        nullptr,               // softmaxLseOptional
        &workspaceSize,        // workspaceSize (out)
        &executor);            // executor (out)

    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBlockSparseAttentionGetWorkspaceSize failed. ERROR: %d\n", ret);
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);
    CHECK_RET(executor != nullptr, LOG_PRINT("executor is null after GetWorkspaceSize\n");
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return -1);

    // 12. Allocate the workspace.
    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);
                      FreeResource(
                      queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor, actualSeqLengths,
                      actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr, valueDeviceAddr,
                      blockSparseMaskDeviceAddr, attentionOutDeviceAddr, actualSeqLengthsDeviceAddr,
                      actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
                  return ret);
    }

    // 12. Call the second-phase API.
    ret = aclnnBlockSparseAttention(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBlockSparseAttention failed. ERROR: %d\n", ret);
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 13. Synchronize and wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret);
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 14. Obtain the output value and copy the result from the device side memory to the host side.
    int64_t attentionOutSize = GetShapeSize(attentionOutShape);
    std::vector<op::fp16_t> resultData(attentionOutSize, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(op::fp16_t), attentionOutDeviceAddr,
                      attentionOutSize * sizeof(op::fp16_t), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret);
              FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor,
                           actualSeqLengths, actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr,
                           valueDeviceAddr, blockSparseMaskDeviceAddr, attentionOutDeviceAddr,
                           actualSeqLengthsDeviceAddr, actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
              return ret);

    // 15. Print partial results.
    uint64_t printNum = 10;
    LOG_PRINT("attentionOut results (first %lu elements):\n", printNum);
    for (uint64_t i = 0; i < printNum && i < resultData.size(); i++) {
        LOG_PRINT("  index %lu: %f\n", i, static_cast<float>(resultData[i]));
    }

    // 16. Free resources.
    FreeResource(queryTensor, keyTensor, valueTensor, blockSparseMaskTensor, attentionOutTensor, actualSeqLengths,
                 actualSeqLengthsKv, blockShape, queryDeviceAddr, keyDeviceAddr, valueDeviceAddr,
                 blockSparseMaskDeviceAddr, attentionOutDeviceAddr, actualSeqLengthsDeviceAddr,
                 actualSeqLengthsKvDeviceAddr, workspaceAddr, deviceId, &stream);
    LOG_PRINT("Test completed successfully!\n");
    return 0;
}