aclnnSparseFlashMlaGrad

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	×
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

• 接口功能：计算SparseFlashMla训练场景下注意力的反向输出，支持Sliding Window Attention、Compressed Attention以及Sparse Compressed Attention。

• 计算公式：

  阶段一：根据不同cmp_ratio场景，对输入ori_kv与cmp_kv进行选择

  • 当cmp_ratio = 1 (SWA)：

  $$selectedKv=orikv$$

  • 当cmp_ratio = 4 (SCFA)：

  $$selectedKv=concat(oriKv,Gather(cmpkv,topkIndices[i])),0<=i<selectBlockCount$$

  • else (CFA):

  $$selectedKv=concat(oriKv,cmpkv)$$

  阶段二：计算P、dP、dS

  $$P=SimpleSoftmax(Mask(Q@selectedK{v}^T\cdot scale),lse)$$

  $$dP=dO@selectedK{v}^T$$

  $$dS=P\times (dP-SoftmaxGrad(dO,O))$$

  阶段三：计算dQ, dKV, dSinks

  $$dQ=dS@selectedKv\cdot scale$$

  $$dKV=d{S}^T@Q\cdot scale+P^T@dO$$

  $$dSinks=ReduceSum(-P\times dP\times SimpleSoftmax(sinks,lse),dim=-1)$$

函数原型

每个算子分为两段式接口，必须先调用“aclnnSparseFlashMlaGradGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnSparseFlashMlaGrad”接口执行计算。

aclnnStatus aclnnSparseFlashMlaGradGetWorkspaceSize(
    const aclTensor   *query,
    const aclTensor   *oriKvOptional,
    const aclTensor   *cmpKvOptional,
    const aclTensor   *dOut,
    const aclTensor   *out,
    const aclTensor   *lse,
    const aclTensor   *oriSparseIndicesOptional,
    const aclTensor   *cmpSparseIndicesOptional,
    const aclTensor   *cuSeqlensQOptional,
    const aclTensor   *cuSeqlensOriKvOptional,
    const aclTensor   *cuSeqlensCmpKvOptional,
    const aclTensor   *sequsedQOptional,
    const aclTensor   *sequsedOriKvOptional,
    const aclTensor   *sequsedCmpKvOptional,
    const aclTensor   *cmpResidualKvOptional,
    const aclTensor   *oriTopkLengthOptional,
    const aclTensor   *cmpTopkLengthOptional,
    const aclTensor   *sinksOptional,
    const aclTensor   *metadataOptional,
    double             scaleValue,
    int64_t            cmpRatio,
    int64_t            oriMaskMode,
    int64_t            cmpMaskMode,
    int64_t            oriWinLeft,
    int64_t            oriWinRight,
    char              *layoutQOptional,
    char              *layoutKvOptional,
    int64_t            deterministic,
    const aclTensor   *dQueryOut,
    const aclTensor   *dOriKvOutOptional,
    const aclTensor   *dCmpKvOutOptional,
    const aclTensor   *dSinksOutOptional,
    const aclTensor   *oriSoftmaxL1NormOptional,
    const aclTensor   *cmpSoftmaxL1NormOptional,
    uint64_t          *workspaceSize,
    aclOpExecutor    **executor);

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

aclnnSparseFlashMlaGradGetWorkspaceSize

• 参数说明：

  参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
  query（aclTensor*）	输入	attention结构的输入Q。	query、oriKvOptional、cmpKvOptional、dOut、outOptional、lse、oriSparseIndicesOptional、cmpSparseIndicesOptional的Shape维度保持一致。	BFLOAT16、FLOAT16	ND	(B,S1,N1,D)、(T1,N1,D) B：支持泛化；S1：支持泛化；N1：支持1~128；D：512；T1：B × S1	√
  oriKvOptional（aclTensor*）	输入	attention结构的原始输入K(V)。	当前暂不支持空tensor。	BFLOAT16、FLOAT16	ND	(B,S2,N2,D)、(T2,N2,D) B：与query的B保持一致；S2：支持泛化；N2：1；D：512；T2：B × S2	√
  cmpKvOptional（aclTensor*）	输入	经过Compressor压缩后的K(V)。	支持空tensor，此时按计算公式中的SWA场景计算。	BFLOAT16、FLOAT16	ND	(B,S3,N2,D)、(T3,N2,D) B：与query的B保持一致；S3 = S2 // cmpRatio；N2：1；D：512；T3：B × S3	√
  dOut（aclTensor*）	输入	注意力输出矩阵的梯度。	-	BFLOAT16、FLOAT16	ND	(B,S1,N1,D)、(T1,N1,D) B：与query的B保持一致；S1：与query的S1保持一致；N1：与query的N1保持一致；D：512；T1：B × S1	√
  out（aclTensor*）	输入	注意力输出矩阵。	-	BFLOAT16、FLOAT16	ND	(B,S1,N1,D)、(T1,N1,D) Shape与dOut保持一致	√
  lse（aclTensor*）	输入	注意力正向计算的输出lse，计算公式详见sparse_flash_mla文档。	-	FLOAT32	ND	(B,N2,S1,G)、(N2,T1,G) B：与query的B保持一致；N2：1；S1：与query的S1保持一致；G：N1/N2；T1：B × S1	√
  oriSparseIndicesOptional（aclTensor*）	输入	稀疏场景下选择的oriKvOptional中权重较高的注意力索引。	目前暂不支持对ori_kv进行稀疏计算，故设置此参数无效。	INT32	ND	(B,S1,N2,K1)、(T1,N2,K1) B：与query的B保持一致；S1：与query的S1保持一致；N2：1；K1：支持泛化；T1：B × S1	√
  cmpSparseIndicesOptional（aclTensor*）	输入	稀疏场景下选择的cmpKvOptional中权重较高的注意力索引。	支持空tensor： 若此时cmpKvOptional非空则按CFA场景计算。 若此时cmpKvOptional为空则按SWA场景计算。	INT32	ND	(B,S1,N2,K2)、(T1,N2,K2) B：与query的B保持一致；S1：与query的S1保持一致；N2：1；K2：支持泛化；T1：B × S1	√
  cuSeqlensQOptional（aclTensor*）	输入	每个Batch中，Query的有效token数。	可选项：当layout为TND，该变量存在。 长度与B+1保持一致。 累加和与T1保持一致。	INT32	ND	(B+1,)	√
  cuSeqlensOriKvOptional（aclTensor*）	输入	每个Batch中，oriKvOptional的有效token数。	可选项：当layout为TND，该变量存在。 长度与B+1保持一致。 累加和与T2保持一致。	INT32	ND	(B+1,)	√
  cuSeqlensCmpKvOptional（aclTensor*）	输入	每个Batch中，cmpKvOptional的有效token数。	可选项：当layout为TND，该变量存在。 长度与B+1保持一致。 累加和与T3保持一致。	INT32	ND	(B+1,)	√
  sequsedQOptional（aclTensor*）	输入	表示不同batch中query实际参与运算的token数。	该参数暂不支持。	INT32	ND	(B,)	√
  sequsedOriKvOptional（aclTensor*）	输入	表示不同batch中oriKvOptional实际参与运算的token数。	该参数暂不支持。	INT32	ND	(B,)	√
  sequsedCmpKvOptional（aclTensor*）	输入	表示不同batch中cmpKvOptional实际参与运算的token数。	该参数暂不支持。	INT32	ND	(B,)	√
  cmpResidualKvOptional（aclTensor*）	输入	表示每个batch S2 // cmpRatio后的余数。	该参数暂不支持。	INT32	ND	(B,)	√
  oriTopkLengthOptional（aclTensor*）	输入	表示每行query对应的oriKvOptional实际可选的topk长度。	该参数暂不支持。	INT32	ND	(B,S1,N2)、(T1,N2)	√
  cmpTopkLengthOptional（aclTensor*）	输入	表示每行query对应的cmpKvOptional实际可选的topk长度。	该参数暂不支持。	INT32	ND	(B,S1,N2)、(T1,N2)	√
  sinksOptional（aclTensor*）	输入	注意力下沉tensor。	-	FLOAT32	ND	(N1) N1：与query的N1保持一致	√
  metadataOptional（aclTensor*）	输入	表示tiling下沉的aicpu算子输出结果。	目前暂不支持tiling下沉，故设置此参数无效。	INT32	ND	(x)	√
  scaleValue（double）	输入	缩放系数。	建议值：公式中d开根号的倒数。

FLOAT32 N/A - - cmpRatio（int64_t） 输入 表示对oriKvOptional的压缩率。 取值范围：1~128。 INT64 N/A - - oriMaskMode（int64_t） 输入 表示query和oriKvOptional计算的mask模式。
• 表示sparse的模式。sparse不同模式的详细说明请参见约束说明。
• 仅支持模式4。
INT64 N/A - - cmpMaskMode（int64_t） 输入 表示query和cmpKvOptional计算的mask模式。
• 表示sparse的模式。sparse不同模式的详细说明请参见约束说明。
• 仅支持模式3。
INT64 N/A - - oriWinLeft（int64_t） 输入 表示query和oriKvOptional计算中query对过去token计算的数量。
• 仅支持取值：127
INT64 - - - oriWinRight（int64_t） 输入 表示query和oriKvOptional计算中query对未来token计算的数量。
• 仅支持取值：0
INT64 - - - layoutQOptional（char*） 输入 表示输入query的数据排布格式。 支持"BSND"、"TND"。 STRING N/A - - layoutKvOptional（char*） 输入 表示输入ori_kv和cmp_kv的数据排布格式。 支持"BSND"、"TND"。 STRING N/A - - deterministic（int64_t） 输入 表示是否开启确定性，应和全局保持一致。 预留参数，暂不支持使用。 INT64 N/A - - dQueryOut（aclTensor*） 输出 表示query的梯度。 与输入query的Shape维度保持一致 BFLOAT16、FLOAT16 ND (B,S1,N1,D)、(T1,N1,D)
√ dOriKvOutOptional（aclTensor*） 输出 表示oriKvOptional的梯度。 与输入oriKvOptional的Shape维度保持一致 BFLOAT16、FLOAT16 ND (B,S2,N2,D)、(T2,N2,D)
√ dCmpKvOptional（aclTensor*） 输出 表示cmpKvOptional的梯度。 与输入cmpKvOptional的Shape维度保持一致。 BFLOAT16、FLOAT16 ND (B,S3,N2,D)、(T3,N2,D)
√ dSinksOutOptional（aclTensor*） 输出 表示sinksOptional的梯度。 与输入sinksOptional的Shape维度保持一致。 FLOAT32 ND (N1)
√ oriSoftmaxL1NormOptional（aclTensor*） 输出 表示query与oriKvOptional计算得出的softmax结果。 该参数暂不支持。 FLOAT32 ND (B,S1,N2,K1)、(T1,N2,K1)
√ cmpSoftmaxL1NormOptional（aclTensor*） 输出 表示query与cmpKvOptional计算得出的softmax结果。 该参数暂不支持。 FLOAT32 ND (B,S1,N2,K2)、(T1,N2,K2)
√

• 返回值：

  返回aclnnStatus状态码，具体参见aclnn返回码。

  第一段接口完成入参校验，出现以下场景时报错：

  返回值	错误码	描述
  ACLNN_ERR_PARAM_NULLPTR	161001	必选参数或者输出是空指针。
  ACLNN_ERR_PARAM_INVALID	161002	输入变量，如query、oriKvOptional、cmpKvOptional……的数据类型和数据格式不在支持的范围内。
  ACLNN_ERR_RUNTIME_ERROR	361001	API内存调用npu runtime的接口异常。

aclnnSparseFlashMlaGrad

• 参数说明：

  参数名	输入/输出	描述
  workspace	输入	在Device侧申请的workspace内存地址。
  workspaceSize	输入	在Device侧申请的workspace大小，由第一段接口aclnnSparseFlashMlaGradGetWorkspaceSize获取。
  executor	输入	op执行器，包含了算子计算流程。
  stream	输入	指定执行任务的Stream流。

• 返回值：

  返回aclnnStatus状态码，具体参见aclnn返回码。

约束说明

• 确定性计算：

  • aclnnSparseFlashMlaGrad默认非确定性实现，暂不支持通过aclrtCtxSetSysParamOpt开启确定性。

• 公共约束

  • 入参为空的场景处理：
    • query为空Tensor：直接返回。

• Mask

  maskMode	含义	备注
  0	不做mask操作	当前不支持
  3	rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。	当前cmp部分支持
  4	band模式的mask，滑窗范围由oriWinLeft、oriWinRight控制，参数起点为右下角。	当前ori部分支持

• 规格约束

  规格项	规格	规格说明
  B	支持泛化	-
  S1、S2	支持泛化	支持S1、S2支持不等长
  N1	1~128	-
  N2	1	-
  D	512	-
  layout	BSND/TND	-

调用示例

调用示例代码如下（以Atlas A2 训练系列产品/Atlas A2 推理系列产品为例），仅供参考，具体编译和执行过程请参考编译与运行样例。

#include <iostream>
#include <vector>
#include <numeric>
#include "acl/acl.h"
#include "aclnnop/aclnn_sparse_flash_mla_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<short> 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: %e\n", i, resultData[i]);
  }
}

int Init(int32_t deviceId, aclrtContext* context, aclrtStream* stream) {
  // 固定写法，AscendCL初始化
  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);
  // 调用aclrtMalloc申请device侧内存
  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);
  // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
  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);

  // 计算连续tensor的strides
  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];
  }

  // 调用aclCreateTensor接口创建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. （固定写法）device/context/stream初始化，参考AscendCL对外接口列表
  // 根据自己的实际device填写deviceId
  int32_t deviceId = 0;
  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. 构造输入与输出，需要根据API的接口自定义构造
  std::vector<int64_t> qShape = {1, 16, 512};                // T1, N1, D
  std::vector<int64_t> oriKvShape = {2048, 1, 512};          // T2, N2, D
  std::vector<int64_t> cmpKvShape = {16, 1, 512};            // T3, N2, D
  std::vector<int64_t> dOutShape = {1, 16, 512};             // T1, N1, D
  std::vector<int64_t> outShape = {1, 16, 512};              // T1, N1, D
  std::vector<int64_t> lseShape = {1, 1, 16};                // N2, T1, G
  std::vector<int64_t> cuSeqQLenshape = {2};                 // B + 1
  std::vector<int64_t> cuSeqOriKvLenshape = {2};             // B + 1
  std::vector<int64_t> cuSeqCmpKvLenshape = {2};             // B + 1
  std::vector<int64_t> sinksShape = {2};                     // N1

  void* qDeviceAddr = nullptr;
  void* oriKvDeviceAddr = nullptr;
  void* cmpKvDeviceAddr = nullptr;
  void* outDeviceAddr = nullptr;
  void* dOutDeviceAddr = nullptr;
  void* lseDeviceAddr = nullptr;
  void* cuSeqQLenDeviceAddr = nullptr;
  void* cuSeqOriKvLenDeviceAddr = nullptr;
  void* cuSeqCmpKvLenDeviceAddr = nullptr;
  void* sinksDeviceAddr = nullptr;
  void* dqDeviceAddr = nullptr;
  void* dOriKvDeviceAddr = nullptr;
  void* dCmpKvDeviceAddr = nullptr;
  void* dSinksDeviceAddr = nullptr;
  void* oriSoftmaxL1NormDeviceAddr = nullptr;
  void* cmpSoftmaxL1NormDeviceAddr = nullptr;

  aclTensor* q = nullptr;
  aclTensor* oriKv = nullptr;
  aclTensor* cmpKv = nullptr;
  aclTensor* out = nullptr;
  aclTensor* dOut = nullptr;
  aclTensor* lse = nullptr;
  aclTensor* cuSeqQLen = nullptr;
  aclTensor* cuSeqOriKvLen = nullptr;
  aclTensor* cuSeqCmpKvLen = nullptr;
  aclTensor* sinks = nullptr;
  aclTensor* dq = nullptr;
  aclTensor* dOriKv = nullptr;
  aclTensor* dCmpKv = nullptr;
  aclTensor* dSinks = nullptr;
  aclTensor* oriSoftmaxL1Norm = nullptr;
  aclTensor* cmpSoftmaxL1Norm = nullptr;

  std::vector<short> qHostData(1 * 16 * 512, 1.0);
  std::vector<short> oriKvHostData(2048 * 1 * 512, 1.0);
  std::vector<short> cmpKvHostData(16 * 1 * 512, 1.0);
  std::vector<short> outHostData(1 * 16 * 512, 1.0);
  std::vector<short> dOutHostData(1 * 16 * 512, 1.0);
  std::vector<float> lseHostData(16, 3.0);
  std::vector<int32_t> cuSeqQLenHostData = {0, 1};
  std::vector<int32_t> cuSeqOriKvLenHostData = {0, 2048};
  std::vector<int32_t> cuSeqCmpKvLenHostData = {0, 16};
  std::vector<float> sinksHostData(16, 128.0);
  std::vector<short> dqHostData(1 * 16 * 512, 0);
  std::vector<short> dOriKvHostData(2048 * 1 * 512, 0);
  std::vector<short> dCmpKvHostData(16 * 1 * 512, 0);
  std::vector<float> dSinksHostData(16, 0);

  ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(oriKvHostData, oriKvShape, &oriKvDeviceAddr, aclDataType::ACL_FLOAT16, &oriKv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(cmpKvHostData, cmpKvShape, &cmpKvDeviceAddr, aclDataType::ACL_FLOAT16, &cmpKv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dOutHostData, dOutShape, &dOutDeviceAddr, aclDataType::ACL_FLOAT16, &dOut);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(lseHostData, lseShape, &lseDeviceAddr, aclDataType::ACL_FLOAT, &lse);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(cuSeqQLenHostData, cuSeqQLenshape, &cuSeqQLenDeviceAddr, aclDataType::ACL_INT32, &cuSeqQLen);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(cuSeqOriKvLenHostData, cuSeqOriKvLenshape, &cuSeqOriKvLenDeviceAddr, aclDataType::ACL_INT32, &cuSeqOriKvLen);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(cuSeqCmpKvLenHostData, cuSeqCmpKvLenshape, &cuSeqCmpKvLenDeviceAddr, aclDataType::ACL_INT32, &cuSeqCmpKvLen);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(sinksHostData, sinksShape, &sinksDeviceAddr, aclDataType::ACL_FLOAT, &sinks);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dqHostData, qShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dOriKvHostData, oriKvShape, &dOriKvDeviceAddr, aclDataType::ACL_FLOAT16, &dOriKv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dCmpKvHostData, cmpKvShape, &dCmpKvDeviceAddr, aclDataType::ACL_FLOAT16, &dCmpKv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(dSinksHostData, sinksShape, &dSinksDeviceAddr, aclDataType::ACL_FLOAT, &dSinks);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  double scaleValue = 0.088388;
  int64_t cmpRatio = 128;
  int64_t oriMaskMode = 4;
  int64_t cmpMaskMode = 3;
  int64_t oriWinLeft = 127;
  int64_t oriWinRight = 0;
  int64_t deterministic = 0;
  char layoutQ[5] = {'T', 'N', 'D', 0};
  char layoutKv[5] = {'T', 'N', 'D', 0};
  
  // 3. 调用CANN算子库API，需要修改为具体的Api名称
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  
  // 调用aclnnSparseFlashMlaGrad第一段接口
  ret = aclnnSparseFlashMlaGradGetWorkspaceSize(q, oriKv, cmpKv, dOut, out, lse, nullptr, nullptr, cuSeqQLen, cuSeqOriKvLen, cuSeqCmpKvLen,
            nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, sinks, nullptr, scaleValue, cmpRatio, oriMaskMode, cmpMaskMode, oriWinLeft, oriWinRight,
            layoutQ, layoutKv, deterministic, dq, dOriKv, dCmpKv, dSinks, oriSoftmaxL1Norm, cmpSoftmaxL1Norm, 
            &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSparseFlashMlaGradGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  
  // 根据第一段接口计算出的workspaceSize申请device内存
  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);
  }
  
  // 调用aclnnSparseFlashMlaGrad第二段接口
  ret = aclnnSparseFlashMlaGrad(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSparseFlashMlaGrad failed. ERROR: %d\n", ret); return ret);
  
  // 4. （固定写法）同步等待任务执行结束
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
  
  // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
  PrintOutResult(qShape, &dqDeviceAddr);
  PrintOutResult(oriKvShape, &dOriKvDeviceAddr);
  PrintOutResult(cmpKvShape, &dCmpKvDeviceAddr);
  PrintOutResult(sinksShape, &dSinksDeviceAddr);

  // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
  aclDestroyTensor(q);
  aclDestroyTensor(oriKv);
  aclDestroyTensor(cmpKv);
  aclDestroyTensor(out);
  aclDestroyTensor(dOut);
  aclDestroyTensor(lse);
  aclDestroyTensor(cuSeqQLen);
  aclDestroyTensor(cuSeqOriKvLen);
  aclDestroyTensor(cuSeqCmpKvLen);
  aclDestroyTensor(sinks);
  aclDestroyTensor(dq);
  aclDestroyTensor(dOriKv);
  aclDestroyTensor(dCmpKv);
  aclDestroyTensor(dSinks);
  aclDestroyTensor(oriSoftmaxL1Norm);
  aclDestroyTensor(cmpSoftmaxL1Norm);

  // 7. 释放device资源
  aclrtFree(qDeviceAddr);
  aclrtFree(oriKvDeviceAddr);
  aclrtFree(cmpKvDeviceAddr);
  aclrtFree(outDeviceAddr);
  aclrtFree(dOutDeviceAddr);
  aclrtFree(lseDeviceAddr);
  aclrtFree(cuSeqQLenDeviceAddr);
  aclrtFree(cuSeqOriKvLenDeviceAddr);
  aclrtFree(cuSeqCmpKvLenDeviceAddr);
  aclrtFree(sinksDeviceAddr);
  aclrtFree(dqDeviceAddr);
  aclrtFree(dOriKvDeviceAddr);
  aclrtFree(dCmpKvDeviceAddr);
  aclrtFree(dSinksDeviceAddr);
  aclrtFree(oriSoftmaxL1NormDeviceAddr);
  aclrtFree(cmpSoftmaxL1NormDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtDestroyContext(context);
  aclrtResetDevice(deviceId);
  aclFinalize();
  
  return 0;
}