aclnnSparseLightningIndexerGradKLLoss

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	√
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Description: The SparselightningIndexerGradKlLoss operator is the backward operator of LightningIndexer, with an additional fused loss computation function. The LightningIndexer operator selects the TopK highest intrinsic connections between QueryToken and KeyToken and stores them in SparseIndices, thereby reducing the computation workload of Attention in long-sequence scenarios and accelerating the inference and training performance of long-sequence networks.
Formula:

The formula for computing the value used to obtain Top-k can be expressed as:
$I_{t,:}=W_{t,:}@ReLU(q_{t,:}@(K_{:t,:})^T)$
Where $W$ is the weights corresponding to the $t$ -th token, $q$ is the matrix formed by concatenating the $G$ query heads corresponding to the $t$ -th token, and $K$ is a matrix with $t$ rows.

LightningIndexer is trained separately, and its corresponding loss function is:
$L(I){=}\sum_tD_{KL}(p_{t,:}||Softmax(I_{t,:}))$
Where $p$ is the target distribution, obtained by summing all heads of the main attention score and then applying L1 normalization to the summation result along the context dimension. $D_{KL}$ is the KL divergence, expressed as:
$D_{KL}(a||b){=}\sum_ia_i\mathrm{log}{\left(\frac{a_i}{b_i}\right)}$
The gradient expression of the loss can be obtained by derivation:
$dI\mathop{{}}\nolimits_{{t,:}}=Softmax \left( I\mathop{{}}\nolimits_{{t,:}} \left) -p\mathop{{}}\nolimits_{{t,:}}\right. \right.$
The gradients of the weights, query, and key matrices can be computed using the chain rule:
$dW\mathop{{}}\nolimits_{{t,:}}=dI\mathop{{}}\nolimits_{{t,:}}\text{@} \left( ReLU \left( S\mathop{{}}\nolimits_{{t,:}} \left) \left) \mathop{{}}\nolimits^{{T}}\right. \right. \right. \right.$ $d\mathop{{q}}\nolimits_{{t,:}}=dS\mathop{{}}\nolimits_{{t,:}}@K\mathop{{}}\nolimits_{{:t,:}}$ $dK\mathop{{}}\nolimits_{{:t,:}}= \left( dS\mathop{{}}\nolimits_{{t,:}} \left) \mathop{{}}\nolimits^{{T}}@q\mathop{{}}\nolimits_{{:t,:}}\right. \right.$
Where S is the result of softmax on the QK matrix.

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the computation process, and then call the "aclnnSparseLightningIndexerGradKLLoss" API to perform the computation.

aclnnStatus aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize(
    const aclTensor     *query,
    const aclTensor     *key,
    const aclTensor     *queryIndex,
    const aclTensor     *keyIndex,
    const aclTensor     *weights,
    const aclTensor     *sparseIndices,
    const aclTensor     *softmaxMax,
    const aclTensor     *softmaxSum,
    const aclTensor     *queryRope,
    const aclTensor     *keyRope,
    const aclIntArray   *actualSeqLengthsQuery,
    const aclIntArray   *actualSeqLengthsKey,
    double               scaleValue,
    char                *layout,
    int64_t              sparseMode,
    int64_t              pre_tokens,
    int64_t              next_tokens,
    bool                 deterministic,
    const aclTensor     *dQueryIndex,
    const aclTensor     *dKeyIndex,
    const aclTensor     *dWeights,
    const aclTensor     *loss,
    uint64_t            *workspaceSize,
    aclOpExecutor       **executor)

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

aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
query	Input	Input Q of the attention structure.	The data type must be consistent with key, queryIndex, and keyIndex.	FLOAT16, BFLOAT16	ND	(B,S1,N1,D), (T1,N1,D)	x
key	Input	Input K of the attention structure.	The data type must be consistent with query, queryIndex, and keyIndex.	FLOAT16, BFLOAT16	ND	(B, S2, N2, D), (T2, N2, D)	x
queryIndex	Input	The input queryIndex of the lightingIndexer structure.	The data type is consistent with query/key/keyIndex.	FLOAT16, BFLOAT16	ND	(B,S1,Nidx1,D), (T1,Nidx1,D)	x
keyIndex	Input	The input keyIndex of the lightingIndexer structure.	The data type must be the same as that of query/key/queryIndex.	FLOAT16, BFLOAT16	ND	(B,S2,Nidx2,D)、(T2,Nidx2,D)	x
weights	Input	Weights.	-	FLOAT16, BFLOAT16, FLOAT32	ND	(B,S1,Nidx1), (T1,Nidx1)	x
sparseIndices	Input	topk_index, used to select the key and value corresponding to each query.	-	INT32	ND	(B,S1,Nidx2,K), (T1,Nidx2,K)	x
softmaxMax	Input	Intermediate output of the attention forward computation.	-	FLOAT32	ND	(B,N2,S1,G)、(N2,T1,G)	x
softmaxSum	Input	Intermediate output of the attention forward computation.	-	FLOAT32	ND	(B,N2,S1,G)、(N2,T1,G)	x
queryRope	Input	MLA RoPE part: Output of the query positional encoding.	Keeps consistent with the layout dimensions of query.	FLOAT16, BFLOAT16	ND	(B,S1,N1,Dr), (T1,N1,Dr)	x
keyRope	Input	Output of the Key position encoding in the MLA rope part.	Consistent with the layout dimension of key.	FLOAT16, BFLOAT16	ND	(B,S2,N2,Dr), (T2,N2,Dr)	x
actualSeqLengthsQuery	Input	Specifies the number of valid tokens in the query for each batch.	Value-dependent. The length must be consistent with B. The cumulative sum must be consistent with T1.	INT64	ND	(B,)	x
actualSeqLengthsKey	Input	Number of valid tokens in Key for each batch.	Value-dependent. The length is consistent with B. The cumulative sum is consistent with T2.	INT64	ND	(B,)	x
scaleValue	Input	Scaling factor.	Recommended value: the reciprocal of the square root of d in the formula.	-	-	-	x
layout	Input	Layout format.	Only BSND and TND formats are supported.	STRING	-	-	x
sparseMode	Input	Sparse mode.	Indicates the sparse mode. For details about different sparse modes, see Constraints. Only mode 3 is supported.	INT64	-	-	x
deterministic	Input	Deterministic computation.	The network-wide deterministic configuration takes precedence. This parameter has no effect.	BOOL	-	-	x
dQueryIndex	Output	Gradient of QueryIndex.	-	FLOAT16, BFLOAT16	ND	(B,S1,Nidx1,D), (T1,Nidx1,D)	x
dKeyIndex	Output	Gradient of KeyIndex.	-	FLOAT16, BFLOAT16	ND	(B,S2,Nidx2,D), (T2,Nidx2,D)	x
dWeights	Output	Gradient of Weights.	-	FLOAT16, BFLOAT16, FLOAT32	ND	(B,S1,Nidx1), (T1,Nidx1)	x
loss	Output	Loss function value.	-	FLOAT32	ND	(1,)	x

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	Return Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	A required parameter or output is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of input variables such as query, key, queryIndex, keyIndex, weights, sparseIndices, and softmaxMax is not within the supported range.
ACLNN_ERR_RUNTIME_ERROR	361001	An exception occurs when the API memory calls the npu runtime interface.

aclnnSparseLightningIndexerGradKLLoss

Parameters:

Parameter	Input/Output	Description
workspace	Input	Address of the workspace memory applied for on the device side.
workspaceSize	Input	Size of the workspace applied for on the device side, obtained through the first-phase API aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize.
executor	Input	Op executor that contains the operator computation process.
stream	Input	Specifies the stream for task execution.

Return value

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

Constraints

Deterministic computation:
- aclnnSparseLightningIndexerGradKLLoss defaults to a non-deterministic implementation and does not support enabling deterministic computation through aclrtCtxSetSysParamOpt.

General Constraints

Handling of empty input parameters:
- query is an empty tensor: return directly.
- The scenario where the input is empty in the common constraints is consistent with FAG.

sparseMode	Meaning	Remarks
0	defaultMask mode. If attenmask is not passed, no mask operation is performed, and preTokens and nextTokens are ignored. If attenmask is passed, the complete attenmask matrix must be passed, indicating that the portion between preTokens and nextTokens needs to be computed.	Not supported.
1	allMask mode. The complete attenmask matrix must be passed.	Not supported.
2	For the leftUpCausal mode mask, the optimized attenmask matrix needs to be passed in.	Not supported.
3	Mask in rightDownCausal mode, corresponding to the lower triangular scenario partitioned by the right vertex. Requires the optimized attenmask matrix.	Supported
4	Mask in band mode. Requires the optimized attenmask matrix.	Not supported
5	prefix	Not supported.
6	global	Not supported.
7	dilated	Not supported.
8	block_local	Not supported.

Specification constraints

Specification Item	Specification	Specification
deterministic	bool	A2/A3 supports deterministic computation. 950 does not support deterministic computation.
B	A2/A3 supports 1 to 256. 950 supports 1 to 128.	-
S1, S2	S1 supports 1 to 8K, and S2 supports 1 to 128K.	S1 and S2 support unequal lengths.
N1	32, 64, 128	SparseFA is MQA.
Nidx1	A2/A3 supports 8, 16, 32, 64. 950 supports 32, 64.	SparseFA is MQA.
N2	1	SparseFA is MQA, and Nidx2 = 1.
Nidx2	1	SparseFA is MQA, with N2=1.
D	512	The D of query differs from that of query_index.
Drope	64	-
K	1024, 2048, 3072, 4096, 5120, 6144, 7168, 8192	-
layout	BSND/TND	-

Typical values

Specification Item	Typical Value
query	N1=128/64; D =512
queryIndex	A2/A3 supports N1 = 64/32/16/8; D = 128; S1 = 64k/128k 950 supports N1 = 64/32; D = 128; S1 = 64k/128k
keyIndex	D = 128
topk	topk = 1024/2048/3072/4096/5120/6144/7168/8192
qRope	d= 64

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 "acl/acl.h"
#include "aclnnop/aclnn_sparse_lightning_indexer_grad_kl_loss.h"

#define CHECK_RET(cond, return_expr) \
  do {                               \
    if (!(cond)) {                   \
      return_expr;                   \
    }                                \
  } while (0)

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
  int64_t shapeSize = 1;
  for (auto i : shape) {
    shapeSize *= i;
  }
  return shapeSize;
}

void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {
  auto size = GetShapeSize(shape);
  std::vector<float> resultData(size, 0);
  auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
                         *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
  for (int64_t i = 0; i < size; i++) {
    LOG_PRINT("mean result[%ld] is: %f\n", i, resultData[i]);
  }
}

int Init(int32_t deviceId, aclrtContext* context, aclrtStream* stream) {
  // 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 = aclrtCreateContext(context, deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetCurrentContext(*context);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor) {
  auto size = GetShapeSize(shape) * sizeof(T);
  // Call aclrtMalloc to apply for device-side memory.
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
  // Copy 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);

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

  // Call the aclCreateTensor API to create an aclTensor.
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

int main() {
  // 1. (Boilerplate) Initialize the device/context/stream. See the AscendCL external API list.
  // Set deviceId based on your actual device.
  int32_t deviceId = 3;
  aclrtContext context;
  aclrtStream stream;
  auto ret = Init(deviceId, &context, &stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct the input and output. Customize the construction based on the API.
  std::vector<int64_t> qShape = {1,64,512};
  std::vector<int64_t> kShape = {1,1,512};
  std::vector<int64_t> qRopeShape = {1,64,64};
  std::vector<int64_t> kRopeShape = {1,1,64};
  std::vector<int64_t> qIndexShape = {1,32,128};
  std::vector<int64_t> kIndexShape = {1,1,128};
  std::vector<int64_t> weightShape = {1,32};
  std::vector<int64_t> sparseIndicesShape = {1, 1,2048};
  std::vector<int64_t> softmaxMaxShape = {1, 1, 64};
  std::vector<int64_t> softmaxSumShape = {1, 1, 64};

  std::vector<int64_t> dQIndexShape = {1,32,128};
  std::vector<int64_t> dKIndexShape = {1,1,128};
  std::vector<int64_t> dWeightShape = {1,1,32};
  std::vector<int64_t> lossShape = {1};

  void* qDeviceAddr = nullptr;
  void* kDeviceAddr = nullptr;
  void* qRopeDeviceAddr = nullptr;
  void* kRopeDeviceAddr = nullptr;
  void* qIndexDeviceAddr = nullptr;
  void* kIndexDeviceAddr = nullptr;
  void* weightDeviceAddr = nullptr;
  void* sparseIndicesDeviceAddr = nullptr;
  void* softmaxMaxDeviceAddr = nullptr;
  void* softmaxSumDeviceAddr = nullptr;

  void* dQIndexDeviceAddr = nullptr;
  void* dKIndexDeviceAddr = nullptr;
  void* dWeightDeviceAddr = nullptr;
  void* lossDeviceAddr = nullptr;

  aclTensor* q = nullptr;
  aclTensor* k = nullptr;
  aclTensor* qRope = nullptr;
  aclTensor* kRope = nullptr;
  aclTensor* qIndex = nullptr;
  aclTensor* kIndex = nullptr;
  aclTensor* weight = nullptr;
  aclTensor* sparseIndices = nullptr;
  aclTensor* softmaxMax = nullptr;
  aclTensor* softmaxSum = nullptr;

  aclTensor* dQIndex = nullptr;
  aclTensor* dKIndex = nullptr;
  aclTensor* dWeight = nullptr;
  aclTensor* loss = nullptr;

  std::vector<float> qHostData(1*64*512, 1);
  std::vector<float> kHostData(1*1*512, 1);
  std::vector<float> qRopeHostData(1*64*64, 1);
  std::vector<float> kRopeHostData(1*1*64, 1);
  std::vector<float> qIndexHostData(1*32*128, 1);
  std::vector<float> kIndexHostData(1*1*128, 1);
  std::vector<float> weightHostData(1*32, 1);
  std::vector<float> sparseIndicesHostData(2048, 1);
  std::vector<float> softmaxMaxHostData(1*64, 1);
  std::vector<float> softmaxSumHostData(1*64, 1);

  std::vector<float> dQIndexHostData(1*32*128, 1);
  std::vector<float> dKIndexHostData(1*1*128, 1);
  std::vector<float> dWeightHostData(1*1*32, 1);
  std::vector<float> lossHostData(1, 1);

  ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(qRopeHostData, qRopeShape, &qRopeDeviceAddr, aclDataType::ACL_FLOAT16, &qRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kRopeHostData, kRopeShape, &kRopeDeviceAddr, aclDataType::ACL_FLOAT16, &kRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(qIndexHostData, qIndexShape, &qIndexDeviceAddr, aclDataType::ACL_FLOAT16, &qIndex);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kIndexHostData, kIndexShape, &kIndexDeviceAddr, aclDataType::ACL_FLOAT16, &kIndex);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr, aclDataType::ACL_FLOAT16, &weight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, &sparseIndicesDeviceAddr, aclDataType::ACL_INT32, &sparseIndices);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  ret = CreateAclTensor(dQIndexHostData, dQIndexShape, &dQIndexDeviceAddr, aclDataType::ACL_FLOAT16, &dQIndex);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dKIndexHostData, dKIndexShape, &dKIndexDeviceAddr, aclDataType::ACL_FLOAT16, &dKIndex);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dWeightHostData, dWeightShape, &dWeightDeviceAddr, aclDataType::ACL_FLOAT16, &dWeight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(lossHostData, lossShape, &lossDeviceAddr, aclDataType::ACL_FLOAT, &loss);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  std::vector<int64_t>  acSeqQLenOp = {1};
  std::vector<int64_t>  acSeqKvLenOp = {1};
  aclIntArray* acSeqQLen = aclCreateIntArray(acSeqQLenOp.data(), acSeqQLenOp.size());
  aclIntArray* acSeqKvLen = aclCreateIntArray(acSeqKvLenOp.data(), acSeqKvLenOp.size());
  double scaleValue = 0.044194173824159216;
  int64_t preTokens = 2147483647;
  int64_t nextTokens = 2147483647;
  int64_t sparseMode = 3;
  bool deterministic = false;

  char layOut[5] = {'T', 'N', 'D', 0};

  // 3. Call the CANN operator library API. Replace with the specific API name.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;

  // Call the aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize first-phase API.
  ret = aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize(
            q, k, qIndex, kIndex, weight, sparseIndices, softmaxMax, softmaxSum, qRope, kRope, acSeqQLen, acSeqKvLen,
            scaleValue, layOut, sparseMode, preTokens, nextTokens, deterministic, dQIndex, dKIndex, dWeight, loss,
            &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSparseLightningIndexerGradKLLossGetWorkspaceSize failed. ERROR: %d\n", ret);
            return ret);

  // Apply for device memory based on the workspaceSize calculated by the first-phase API.
  void* workspaceAddr = nullptr;
  if (workspaceSize > 0) {
    ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
  }

  // Call the aclnnSparseLightningIndexerGradKLLoss second-phase API.
  ret = aclnnSparseLightningIndexerGradKLLoss(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSparseLightningIndexerGradKLLoss failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait for the task execution to finish synchronously.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 5. Obtain the output value and copy the result from the device memory to the host memory. Modify this based on the specific API definition.
  PrintOutResult(dQIndexShape, &dQIndexDeviceAddr);
  PrintOutResult(dKIndexShape, &dKIndexDeviceAddr);
  PrintOutResult(dWeightShape, &dWeightDeviceAddr);
  PrintOutResult(lossShape, &lossDeviceAddr);

  // 6. Release aclTensor and aclScalar. Modify this based on the specific API definition.
  aclDestroyTensor(q);
  aclDestroyTensor(k);
  aclDestroyTensor(qIndex);
  aclDestroyTensor(kIndex);
  aclDestroyTensor(qRope);
  aclDestroyTensor(kRope);
  aclDestroyTensor(weight);
  aclDestroyTensor(sparseIndices);
  aclDestroyTensor(softmaxMax);
  aclDestroyTensor(softmaxSum);

  aclDestroyTensor(dQIndex);
  aclDestroyTensor(dKIndex);
  aclDestroyTensor(dWeight);
  aclDestroyTensor(loss);

  // 7. Release device resources.
  aclrtFree(qDeviceAddr);
  aclrtFree(kDeviceAddr);
  aclrtFree(qIndexDeviceAddr);
  aclrtFree(kIndexDeviceAddr);
  aclrtFree(qRopeDeviceAddr);
  aclrtFree(kRopeDeviceAddr);
  aclrtFree(weightDeviceAddr);
  aclrtFree(sparseIndicesDeviceAddr);
  aclrtFree(softmaxMaxDeviceAddr);
  aclrtFree(softmaxSumDeviceAddr);

  aclrtFree(dQIndexDeviceAddr);
  aclrtFree(dKIndexDeviceAddr);
  aclrtFree(dWeightDeviceAddr);
  aclrtFree(lossDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtDestroyContext(context);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}