aclnnFlashAttentionUnpaddingScoreGradV3

Applicable Products

Function

• Description: Computes the backward output of attention in training scenarios, i.e., the backward computation of aclnnFlashAttentionVarLenScoreV3. Compared with the aclnnFlashAttentionUnpaddingScoreGradV2 API, this API adds the queryRope, keyRope, dqRope, and dkRope parameters.

• Formula:

  The forward formula of attention is as follows:

  $$Y=Dropout(Softmax(Mask(\frac{QK^T}{\sqrt{d}}+pse),atten\_mask),keep\_prob)V$$

  Where:

  $$Q=[query, queryRope], K=[key, keyRope]$$

  For ease of expression, the variables $S$ and $P$ are used to represent the formulas:

  $$S=Mask(\frac{QK^T}{\sqrt{d}}+pse),atten\_mask$$ $$P=Dropout(Softmax(S),keep\_prob)$$ $$Y=PV$$

  The backward calculation formula for attention is:

  $$dV=P^TdY$$ $$dQ=\frac{((dS)*K)}{\sqrt{d}}$$ $$dqRope=\frac{((dS)*kRope)}{\sqrt{d}}$$ $$dK=\frac{((dS)^T*Q)}{\sqrt{d}}$$ $$dkRope=\frac{((dS)^T*qRope)}{\sqrt{d}}$$

Function Prototype

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

aclnnStatus aclnnFlashAttentionUnpaddingScoreGradV3GetWorkspaceSize(
  const aclTensor   *query,
  const aclTensor   *queryRope,
  const aclTensor   *keyIn,
  const aclTensor   *keyInRope,
  const aclTensor   *value,
  const aclTensor   *dy,
  const aclTensor   *pseShiftOptional,
  const aclTensor   *dropMaskOptional,
  const aclTensor   *paddingMaskOptional,
  const aclTensor   *attenMaskOptional,
  const aclTensor   *softmaxMaxOptional,
  const aclTensor   *softmaxSumOptional,
  const aclTensor   *softmaxInOptional,
  const aclTensor   *attentionInOptional,
  const aclIntArray *prefixOptional,
  const aclIntArray *actualSeqQLenOptional,
  const aclIntArray *actualSeqKvLenOptional,
  const aclIntArray *qStartIdxOptional,
  const aclIntArray *kvStartIdxOptional,
  double             scaleValue,
  double             keepProb,
  int64_t            preTokens,
  int64_t            nextTokens,
  int64_t            headNum,
  char              *inputLayout,
  int64_t            innerPrecise,
  int64_t            sparseMode,
  int64_t            pseType,
  const aclTensor   *dqOut,
  const aclTensor   *dqRopeOut,
  const aclTensor   *dkOut,
  const aclTensor   *dkRopeOut,
  const aclTensor   *dvOut,
  const aclTensor   *dpseOut,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

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

aclnnFlashAttentionUnpaddingScoreGradV3GetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
  query	Input.	Q in the formula.	The data type is consistent with keyIn/value.	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  queryRope	optional input	The rope part (rotary position encoding) of Q.	Data type consistent with query.	BFLOAT16	ND	[TND]	√
  keyIn	Input	K in the formula.	The data type is consistent with query and value.	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  keyInRope	optional input	The rope part of K.	The data type is consistent with keyIn.	BFLOAT16	ND	[TND]	√
  value	Input.	V in the formula.	The data type is consistent with query/keyIn.	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	Yes.
  dy	Input	dY in the formula.	-	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  pseShiftOptional	optional input	pse in the formula, the positional encoding.	Must be nullptr.	-	-	-	-
  dropMaskOptional	Optional input	Dropout mask.	Must be nullptr.	-	-	-	-
  paddingMaskOptional	optional input	Reserved parameter, not currently used.	Pass null when calling.	-	-	-	-
  qStartIdxOptional	Optional input	Indicates the start index of the current chunked Query in the global sequence in an external slicing scenario.	-	INT64	ND	0, 1	-
  kvStartIdxOptional	Optional input	Start index of the current chunk Key/Value in the global sequence in the external slicing scenario.	-	INT64	ND	0, 1	-
  attenMaskOptional	Optional input	Represents atten_mask in the formula.	A value of 1 indicates that the bit does not participate in the computation, and 0 indicates that the bit participates in the computation.	BOOL, UINT8	ND	[B,N,Sq,Skv], [B,1,Sq,Skv], [1,1,Sq,Skv], [Sq,Skv]	√
  softmaxMaxOptional	Optional input	Intermediate output of the forward softmax.	-	FLOAT	ND	[N,T,8]	√
  softmaxSumOptional	Optional input.	Intermediate output of the forward softmax.	-	FLOAT	ND	[N,T,8]	√
  softmaxInOptional	Optional input	Intermediate output of the forward softmax.	Reserved parameter, not currently used.	-	-	-	-
  attentionInOptional	Optional input.	Forward attention output.	Must have the same data type and shape as query.	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  prefixOptional	Optional input	N value of each batch in the prefix sparse scenario.	nullptr can be passed.	INT64	ND	0, 1	-
  actualSeqQLenOptional	Optional input	Specifies the actual Query sequence length.	-	INT64	ND	1	-
  actualSeqKvLenOptional	Optional input	Actual key/value sequence length.	-	INT64	ND	1	-
  scaleValue	Optional input	Scale factor.	-	DOUBLE	-	-	-
  keepProb	Optional Input	Proportion of 1s in dropMaskOptional.	-	DOUBLE	-	-	-
  preTokens	optional input	Left boundary of the sliding window in sparse computation.	-	INT64	-	-	-
  nextTokens	Optional input	Specifies the right boundary of the sliding window during sparse computation.	-	INT64	-	-	-
  headNum	Input	Number of heads on a single device, i.e., the length of the N-axis of Query.	-	INT64	-	-	-
  inputLayout	Input	Specifies the layout of the input Q/K/V data.	Supports TND.	String	-	-	-
  innerPrecise	Optional input	Internal computation precision control.	Reserved parameter, not currently used.	INT64	-	-	-
  sparseMode	Optional input	Sparse mode.	Supports configuration from 0 to 8. 5 and 6 are not supported.	INT64	-	-	-
  pseType	Optional input	Supports the INT64 data type.	Only the value 1 is supported.	INT64	-	-	-
  dqOut	Output	dQ in the formula, the gradient of Query.	-	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  dqRopeOut	Output	Rope gradient of Q.	-	BFLOAT16	ND	[TND]	√
  dkOut	Output	dK in the formula, representing the gradient of Key.	-	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  dkRopeOut	Output	Gradient of K's rope.	-	BFLOAT16	ND	[TND]	√
  dvOut	Output	dV in the formula, representing the value gradient.	-	FLOAT16, BFLOAT16, FLOAT32	ND	[TND]	√
  dpseOut	Output	Gradient of d(pse).	Reserved parameter, not currently used.	-	-	-	-
  workspaceSize	Output	Returns the workspace size to be side-applied on the Device.	-	-	-	-	-
  executor	Output	Returns the operator executor, which contains 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	A required input, output, or required attribute is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data types of query, queryRope, keyIn, keyInRope, value, dy, pseShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOptional, softmaxSumOptional, softmaxInOptional, attentionInOptional, dqOut, dkOut, and dvOut are not within the supported range.
  		The data formats of query, queryRope, keyIn, keyInRope, value, dy, pseShiftOptional, dropMaskOptional, paddingMaskOptional, attenMaskOptional, softmaxMaxOptional, softmaxSumOptional, softmaxInOptional, attentionInOptional, dqOut, dkOut, and dvOut are not within the supported range.

aclnnFlashAttentionUnpaddingScoreGradV3

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Device-side workspace memory address.
  workspaceSize	Input	Device-side workspace size, obtained from the first-phase API aclnnFlashAttentionUnpaddingScoreGradV3GetWorkspaceSize.
  executor	Input	Operator executor, which contains the operator computation flow.
  stream	Input	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic computation:

• aclnnFlashAttentionUnpaddingScoreGradV3 defaults to a non-deterministic implementation. Deterministic computation can be enabled through aclrtCtxSetSysParamOpt.

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

• The B (batch size) of the input query, queryRope, key, keyRope, value, and dy must be equal.

• The inputLayout of the input query, queryRope, key, keyRope, value, and dy must be TND.

• When the d dimensions of query, key, and value are the same, the shapes of query and dy must be consistent.

• The d dimensions of query and key must be the same, and d must be an integer multiple of 8.

• The d dimensions of queryRope and keyRope must be the same, and d must be an integer multiple of 8 and must be less than or equal to the d of query and key.

• The N of query and dy and the N of key and value can be different, but they must be in a proportional relationship, that is, Nq/Nkv must be a non-zero integer, with Nq in the value range of 1 to 256.

• Constraints on data shape, taking TND of inputLayout as an example:

  • T: value range is 1 to 1M.

  • N: value range is 1 to 256.

  • D: value range is 1 to 768.

  • KeepProb: value range is 1.

• The data layout of query, key, and value supports only TND, where T is the tightly packed data of B and S axes (SeqLenQ and SeqLenKV of each batch), B (Batch) indicates the input sample batch size, S (Seq-Length) indicates the input sample sequence length, H (Head-Size) indicates the hidden size, N (Head-Num) indicates the number of multi-heads, and D (Head-Dim) indicates the minimum unit size of the hidden layer, satisfying D = H/N.

• Constraints on sparseMode are as follows:

  • When all attenMaskOptional shapes are smaller than 2048 and identical, it is recommended to use the default mode to reduce memory usage.

  • When configured as 1, 2, or 3, the user-configured preTokens and nextTokens do not take effect.

  • When configured as 0 or 4, ensure that attenMaskOptional is consistent with the ranges of preTokens and nextTokens.

  • It is recommended to pass 0 if not specifically specified by the user.

  • For details about different sparse modes, see sparse mode description.

  • When configured as 7, the optional input pseShiftOptional is not supported.

  • When configured as 8, the optional input pseShiftOptional is supported when q and kv of each sequence have equal lengths, and pse is generated globally. External slicing in the q direction is supported, which requires that q and kv of each sequence have equal lengths before external slicing, and after external slicing, actualSeqQLenOptional[0] - actualSeqKvLenOptional[0] + qStartIdxOptional - kvStartIdxOptional == 0 (this feature is experimental).

• In some scenarios, if the computation workload is too large, the operator execution may time out (aicore error type, with errorStr: timeout or trap error). In this case, axis splitting is recommended. Note: The computation workload here is affected by parameters such as B, S, N, and D. Larger values result in greater computation workload.

• actualSeqQLenOptional input supports an S length of 0 in a batch, but the optional input pseShiftOptional is not supported in this case. The length of actualSeqQLenOptional ranges from 1 to 2K. When the prefixOptional input is present, its maximum length is 1K.

• Constraints on the softmaxMax and softmaxSum parameters: The input format is fixed as [T, N, 8], where T = B * S.

• The value of headNum must be consistent with the N value in the input Query.

• pseType can only be 1.

• pseShiftOptional must be empty.

• dropMaskOptional must be empty.

• attenMaskOptional cannot be empty.

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_flash_attention_score_grad.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, aclrtStream* stream) {
  // Boilerplate: initialize resources.
  auto ret = aclInit(nullptr);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetDevice(deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  return 0;
}

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

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

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

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

  // 2. Construct the input and output. Customize the construction based on the API interface.
  std::vector<int64_t> qShape = {256, 1, 128};
  std::vector<int64_t> qRopeShape = {256, 1, 64};
  std::vector<int64_t> kShape = {256, 1, 128};
  std::vector<int64_t> kRopeShape = {256, 1, 64};
  std::vector<int64_t> vShape = {256, 1, 128};
  std::vector<int64_t> dxShape = {256, 1, 128};
  std::vector<int64_t> attenmaskShape = {256, 256};
  std::vector<int64_t> softmaxMaxShape = {256, 1, 8};
  std::vector<int64_t> softmaxSumShape = {256, 1, 8};
  std::vector<int64_t> attentionInShape = {256, 1, 128};

  std::vector<int64_t> dqShape = {256, 1, 128};
  std::vector<int64_t> dqRopeShape = {256, 1, 64};
  std::vector<int64_t> dkShape = {256, 1, 128};
  std::vector<int64_t> dkRopeShape = {256, 1, 64};
  std::vector<int64_t> dvShape = {256, 1, 128};

  void* qDeviceAddr = nullptr;
  void* qRopeDeviceAddr = nullptr;
  void* kDeviceAddr = nullptr;
  void* kRopeDeviceAddr = nullptr;
  void* vDeviceAddr = nullptr;
  void* dxDeviceAddr = nullptr;
  void* attenmaskDeviceAddr = nullptr;
  void* softmaxMaxDeviceAddr = nullptr;
  void* softmaxSumDeviceAddr = nullptr;
  void* attentionInDeviceAddr = nullptr;
  void* dqDeviceAddr = nullptr;
  void* dqRopeDeviceAddr = nullptr;
  void* dkDeviceAddr = nullptr;
  void* dkRopeDeviceAddr = nullptr;
  void* dvDeviceAddr = nullptr;

  aclTensor* q = nullptr;
  aclTensor* qRope = nullptr;
  aclTensor* k = nullptr;
  aclTensor* kRope = nullptr;
  aclTensor* v = nullptr;
  aclTensor* dx = nullptr;
  aclTensor* pse = nullptr;
  aclTensor* dropMask = nullptr;
  aclTensor* padding = nullptr;
  aclTensor* attenmask = nullptr;
  aclTensor* softmaxMax = nullptr;
  aclTensor* softmaxSum = nullptr;
  aclTensor* softmaxIn = nullptr;
  aclTensor* attentionIn = nullptr;
  aclTensor* dq = nullptr;
  aclTensor* dqRope = nullptr;
  aclTensor* dk = nullptr;
  aclTensor* dkRope = nullptr;
  aclTensor* dv = nullptr;
  aclTensor* dpse = nullptr;

  std::vector<float> qHostData(32768, 1);
  std::vector<float> qRopeHostData(16384, 1);
  std::vector<float> kHostData(32768, 1);
  std::vector<float> kRopeHostData(16384, 1);
  std::vector<float> vHostData(32768, 1);
  std::vector<float> dxHostData(32768, 1);
  std::vector<uint8_t> attenmaskHostData(65536, 0);
  std::vector<float> softmaxMaxHostData(2048, 3.0);
  std::vector<float> softmaxSumHostData(2048, 3.0);
  std::vector<float> attentionInHostData(32768, 1);
  std::vector<float> dqHostData(32768, 0);
  std::vector<float> dqRopeHostData(16384, 0);
  std::vector<float> dkHostData(32768, 0);
  std::vector<float> dkRopeHostData(16384, 0);
  std::vector<float> dvHostData(32768, 0);

  ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_BF16, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(qRopeHostData, qRopeShape, &qRopeDeviceAddr, aclDataType::ACL_BF16, &qRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_BF16, &k);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(kRopeHostData, kRopeShape, &kRopeDeviceAddr, aclDataType::ACL_BF16, &kRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_BF16, &v);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_BF16, &dx);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
  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(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_BF16, &attentionIn);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_BF16, &dq);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dqRopeHostData, dqRopeShape, &dqRopeDeviceAddr, aclDataType::ACL_BF16, &dqRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_BF16, &dk);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dkRopeHostData, dkRopeShape, &dkRopeDeviceAddr, aclDataType::ACL_BF16, &dkRope);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_BF16, &dv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  std::vector<int64_t> prefixOp = {0};
  aclIntArray* prefix = aclCreateIntArray(prefixOp.data(), 1);
  std::vector<int64_t>  acSeqQLenOp = {256};
  std::vector<int64_t>  acSeqKvLenOp = {256};
  aclIntArray* acSeqQLen = aclCreateIntArray(acSeqQLenOp.data(), acSeqQLenOp.size());
  aclIntArray* acSeqKvLen = aclCreateIntArray(acSeqKvLenOp.data(), acSeqKvLenOp.size());
  std::vector<int64_t> qStartIdxOp = {0};
  std::vector<int64_t> kvStartIdxOp = {0};
  aclIntArray *qStartIdx = aclCreateIntArray(qStartIdxOp.data(), 1);
  aclIntArray *kvStartIdx = aclCreateIntArray(kvStartIdxOp.data(), 1);
  double scaleValue = 0.088388;
  double keepProb = 1;
  int64_t preTokens = 65536;
  int64_t nextTokens = 65536;
  int64_t headNum = 1;
  int64_t innerPrecise = 0;
  int64_t sparseMode = 0;
  int64_t pseType = 1;
  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 first-phase API of aclnnFlashAttentionUnpaddingScoreGradV3.
  ret = aclnnFlashAttentionUnpaddingScoreGradV3GetWorkspaceSize(q, qRope, k, kRope, v, dx, pse, dropMask, padding,
            attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, prefix, acSeqQLen, acSeqKvLen, qStartIdx, kvStartIdx,
            scaleValue, keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise, sparseMode, pseType,
            dq, dqRope, dk, dkRope, dv, dpse, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionUnpaddingScoreGradV3GetWorkspaceSize 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 second-phase API of aclnnFlashAttentionUnpaddingScoreGradV3.
  ret = aclnnFlashAttentionUnpaddingScoreGradV3(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionUnpaddingScoreGradV3 failed. ERROR: %d\n", ret); return ret);

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

  // 5. Obtain the output value and copy the result from device memory to host memory. Modify based on the specific API definition.
  PrintOutResult(dqShape, &dqDeviceAddr);
  PrintOutResult(dqRopeShape, &dqRopeDeviceAddr);
  PrintOutResult(dkShape, &dkDeviceAddr);
  PrintOutResult(dkRopeShape, &dkRopeDeviceAddr);
  PrintOutResult(dvShape, &dvDeviceAddr);

  // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
  aclDestroyTensor(q);
  aclDestroyTensor(qRope);
  aclDestroyTensor(k);
  aclDestroyTensor(kRope);
  aclDestroyTensor(v);
  aclDestroyTensor(dx);
  aclDestroyTensor(attenmask);
  aclDestroyTensor(softmaxMax);
  aclDestroyTensor(softmaxSum);
  aclDestroyTensor(attentionIn);
  aclDestroyTensor(dq);
  aclDestroyTensor(dqRope);
  aclDestroyTensor(dk);
  aclDestroyTensor(dkRope);
  aclDestroyTensor(dv);

  // 7. Release device resources.
  aclrtFree(qDeviceAddr);
  aclrtFree(qRopeDeviceAddr);
  aclrtFree(kDeviceAddr);
  aclrtFree(kRopeDeviceAddr);
  aclrtFree(vDeviceAddr);
  aclrtFree(dxDeviceAddr);
  aclrtFree(attenmaskDeviceAddr);
  aclrtFree(softmaxMaxDeviceAddr);
  aclrtFree(softmaxSumDeviceAddr);
  aclrtFree(attentionInDeviceAddr);
  aclrtFree(dqDeviceAddr);
  aclrtFree(dqRopeDeviceAddr);
  aclrtFree(dkDeviceAddr);
  aclrtFree(dkRopeDeviceAddr);
  aclrtFree(dvDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();

  return 0;
}