aclnnSparseFlashAttentionV2【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer 查看源码产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明接口功能sparse_flash_attentionSFA是针对大序列长度推理场景的高效注意力计算模块该模块通过“只计算关键部分”大幅减少计算量然而会引入大量的离散访存造成数据搬运时间增加进而影响整体性能。V2版本新增sinks参数。计算公式$$ \text{softmax}(\frac{Q\tilde{K}^T}{\sqrt{d_k}})\tilde{V} $$其中$\tilde{K},\tilde{V}$为基于某种选择算法如lightning_indexer得到的重要性较高的Key和Value一般具有稀疏或分块稀疏的特征$d_k$为$Q,\tilde{K}$每一个头的维度。函数原型每个算子分为两段式接口必须先调用“aclnnSparseFlashAttentionV2GetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器再调用“aclnnSparseFlashAttentionV2”接口执行计算。aclnnStatus aclnnSparseFlashAttentionV2GetWorkspaceSize( const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *sparseIndices, const aclTensor *blockTableOptional, const aclTensor *actualSeqLengthsQueryOptional, const aclTensor *actualSeqLengthsKvOptional, const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, const aclTensor *sinksOptional, double scaleValue, int64_t sparseBlockSizeOptional, char *layoutQueryOptional, char *layoutKvOptional, int64_t sparseMode, int64_t preTokens, int64_t nextTokens, int64_t attentionMode, bool returnSoftmaxLse, const aclTensor *attentionOutOut, const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnSparseFlashAttentionV2( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)aclnnSparseFlashAttentionV2GetWorkspaceSize参数说明[!NOTE]query、key、value参数维度含义BBatch Size表示输入样本批量大小、SSequence Length表示输入样本序列长度、HHead Size表示hidden层的大小、NHead Num表示多头数、DHead Dim表示hidden层最小的单元尺寸且满足DH/N、T表示所有Batch输入样本序列长度的累加和。Q_S和S1表示query shape中的SKV_S和S2表示key shape中的SQ_N和N1表示num_query_headsKV_N和N2表示num_key_value_headsT1表示query shape中的TT2表示key shape中的输入样本序列长度的累加和。参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorqueryaclTensor输入attention结构的Query输入。不支持空tensor。FLOAT16、BFLOAT16NDlayout_query为BSND时shape为(B,S1,N1,D)。layout_query为TND时shape为(T1,N1,D)。xkeyaclTensor输入attention结构的Key输入不支持空tensor。block_num为PageAttention时block总数。FLOAT16、BFLOAT16NDlayout_kv为PA_BSND时shape为(block_num, block_size, KV_N, D)。layout_kv为BSND时shape为(B, S2, KV_N, D)。layout_kv为TND时shape为(T2, KV_N, D)。xvalueaclTensor输入attention结构的Value输入。不支持空tensor。FLOAT16、BFLOAT16NDshape与key的shape一致。xsparseIndicesaclTensor输入离散取kvCache的索引。不支持空tensor。sparse_size为一次离散选取的block数需要保证每行有效值均在前半部分无效值均在后半部分且需要满足sparse_size大于0。INT32NDlayout_query为BSND时shape为(B, Q_S, KV_N, sparse_size)。layout_query为TND时shape为(Q_T, KV_N, sparse_size)。xblockTableOptionalaclTensor输入表示PageAttention中kvCache存储使用的block映射表。不支持空tensor。第二维长度不小于所有batch中最大的S2对应的block数量即S2_max / block_size向上取整。INT32NDshape支持(B,S2/block_size)。xactualSeqLengthsQueryOptionalaclTensor输入表示不同Batch中query的有效token数。不支持空tensor。如果不指定seqlen可传入None表示和query的shape的S长度相同。该入参中每个Batch的有效token数不超过query中的维度S大小且不小于0。支持长度为B的一维tensor。layout_query为TND时该入参必须传入且以该入参元素的数量作为B值该参数中每个元素的值表示当前batch与之前所有batch的token数总和。INT32ND(B,)xactualSeqLengthsKvOptionalaclTensor输入表示不同Batch中key和value的有效token数。不支持空tensor。如果不指定seqlen可传入None表示和key的shape的S长度相同。该参数中每个Batch的有效token数不超过key/value中的维度S大小且不小于0。支持长度为B的一维tensor。当layout_kv为TND或PA_BSND时该入参必须传入。layout_kv为TND该参数中每个元素的值表示当前batch与之前所有batch的token数总和即前缀和因此后一个元素的值必须大于等于前一个元素的值。INT32ND(B,)xqueryRopeOptionalaclTensor输入表示MLA结构中的query的rope信息。不支持空tensor。FLOAT16、BFLOAT16NDlayout_query为TND时shape为(B,S1,N1,Dr)。layout_query为BSND时shape为(T1,N1,Dr)。xkeyRopeOptionalaclTensor输入表示MLA结构中的key的rope信息。不支持空tensor。FLOAT16、BFLOAT16NDlayout_kv为TND时shape为(B,S1,N1,Dr)。layout_kv为BSND时shape为(T1,N1,Dr)。layout_kv为PA_BSND时shape为(block_num,block_size,N2,Dr)。xsinksOptionalaclTensor输入attention结构中的可学习的sinks信息。不支持空tensor。FLOATNDshape支持(N1)。xscaleValuedouble输入代表缩放系数。-FLOAT16---sparseBlockSizeOptionalint64_t输入代表sparse阶段的block大小。sparse_block_size为1时为Token-wise稀疏化场景将每个token视为独立单元在计算重要性分数时评估每个查询token与每个键值token之间的独立关联程度。sparse_block_size为大于1小于等于128时为Block-wise稀疏化场景将token序列划分为固定大小的连续块以块为单位进行重要性评估块内token共享相同的稀疏化决策。INT64---layoutQueryOptionalchar输入标识输入query的数据排布格式。用户不特意指定时可传入默认值BSND。支持传入BSND和TND。STRING---layoutKvOptionalchar输入标识输入key的数据排布格式。用户不特意指定时可传入默认值BSND。支持传入TND、BSND和PA_BSND其中PA_BSND在开启PageAttention时使用。STRING---sparseModeint64_t输入表示sparse的模式。sparse_mode为0时代表全部计算。sparse_mode为3时代表rightDownCausal模式的mask对应以右下顶点往左上为划分线的下三角场景。INT64---preTokensint64_t输入用于稀疏计算表示attention需要和前几个Token计算关联。仅支持默认值2^63-1。INT64---nextTokensint64_t输入用于稀疏计算表示attention需要和后几个Token计算关联。仅支持默认值2^63-1。INT64---attentionModeint64_t输入-仅支持传入2表示MLA-absorb模式。INT64---returnSoftmaxLsebool输入用于表示是否返回softmax_max和softmax_sum。True表示返回False表示不返回默认值为False。该参数仅在训练且layout_kv不为PA_BSND场景支持。BOOL---attentionOutaclTensor输出公式中的输出。不支持空tensor。FLOAT16、BFLOAT16NDlayout_query为BSND时shape为(B,S1,N1,D)。layout_query为TND时shape为(T1,N1,D)。xsoftmaxMaxOutaclTensor输出Attention算法对query乘key的结果取max得到softmax_max。不支持空tensor。FLOATNDlayout_query为BSND时shape为(B,N2,S1,N1/N2)。layout_query为TND时shape为(N2,T1,N1/N2)。xsoftmaxSumOutaclTensor输出Attention算法query乘key的结果减去softmax_max,再取exp接着求sum得到softmax_sum。不支持空tensor。FLOATNDlayout_query为BSND时shape为(B,N2,S1,N1/N2)。layout_query为TND时shape为(N2,T1,N1/N2)。xworkspaceSizeuint64_t*输出返回需要在Device侧申请的workspace大小。-----executoraclOpExecutor输出返回op执行器包含了算子计算流程。-----返回值aclnnStatus返回状态码具体参见aclnn返回码。第一段接口会完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入输出或者必选属性且是空指针则返回161001。ACLNN_ERR_PARAM_INVALID161002query、key、value、sparseIndices、blockTableOptional、actualSeqLengthsQueryOptional、actualSeqLengthsKvOptional、queryRopeOptional、keyRopeOptional、sinksOptional、scaleValue、sparseBlockSizeOptional、layoutQueryOptional、layoutKvOptional、sparseMode、attentionMode、returnSoftmaxLse、attentionOut、softmaxMaxOut、softmaxSumOut的数据类型和数据格式不在支持的范围内。aclnnSparseFlashAttentionV2参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnSparseFlashAttentionV2GetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的Stream。返回值aclnnStatus返回状态码具体参见aclnn返回码。约束说明确定性计算aclnnSparseFlashAttentionV2默认确定性实现。该接口支持推理场景下使用。N1支持情况Ascend 950PR/Ascend 950DT N1支持1~128。Atlas A3 训练系列产品/Atlas A3 推理系列产品 、 Atlas A2 训练系列产品/Atlas A2 推理系列产品 N1支持1/2/4/8/16/32/64/128。N2仅支持1。block_size为一个block的token数block_size取值为16的倍数且最大支持1024。参数query中的D和key、value的D值相等为512参数query_rope中的Dr和key_rope的Dr值相等为64。参数query、key、value的数据类型必须保持一致。支持sparse_block_size整除block_size。Ascend 950PR/Ascend 950DT 只支持sparse_block_size为1。Atlas A3 训练系列产品/Atlas A3 推理系列产品 、 Atlas A2 训练系列产品/Atlas A2 推理系列产品 支持[1,128]且要求是2的幂次方在PageAttention场景下要求sparse_block_size整除block_size参数sinks仅支持Ascend 950PR/Ascend 950DT。调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。/** * Copyright (c) 2024 Huawei Technologies Co., Ltd. * This file is a part of the CANN Open Software. * Licensed under CANN Open Software License Agreement Version 1.0 (the License). * Please refer to the License for details. You may not use this file except in compliance with the License. * THIS SOFTWARE IS PROVIDED ON AN AS IS BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. * See LICENSE in the root of the software repository for the full text of the License. */ /*! * \file test_incre_flash_attention_v4.cpp * \brief */ #include iostream #include vector #include cmath #include cstring #include securec.h #include acl/acl.h #include aclnnop/aclnn_sparse_flash_attention.h using namespace std; namespace { #define CHECK_RET(cond) ((cond) ? true :(false)) #define LOG_PRINT(message, ...) \ do { \ (void)printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret; } ret aclrtSetDevice(deviceId); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret; } ret aclrtCreateStream(stream); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret; } return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret; } ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret; } std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } struct TensorResources { void* queryDeviceAddr nullptr; void* keyDeviceAddr nullptr; void* valueDeviceAddr nullptr; void* sparseIndicesDeviceAddr nullptr; void* attentionOutDeviceAddr nullptr; void* softmaxMaxDeviceAddr nullptr; void* softmaxSumDeviceAddr nullptr; void* queryRopeDeviceAddr nullptr; void* keyRopeDeviceAddr nullptr; void* sinksDeviceAddr nullptr; aclTensor* queryTensor nullptr; aclTensor* keyTensor nullptr; aclTensor* valueTensor nullptr; aclTensor* sparseIndicesTensor nullptr; aclTensor* attentionOutTensor nullptr; aclTensor* softmaxMaxTensor nullptr; aclTensor* softmaxSumTensor nullptr; aclTensor* queryRopeTensor nullptr; aclTensor* keyRopeTensor nullptr; aclTensor* sinksTensor nullptr; }; int InitializeTensors(TensorResources resources) { std::vectorint64_t queryShape {1, 2, 1, 512}; std::vectorint64_t keyShape {1, 2, 1, 512}; std::vectorint64_t valueShape {1, 2, 1, 512}; std::vectorint64_t sparseIndicesShape {1, 2, 1, 2}; std::vectorint64_t attentionOutShape {1, 2, 1, 512}; std::vectorint64_t softmaxMaxShape {1, 2, 1, 16}; std::vectorint64_t softmaxSumShape {1, 2, 1, 16}; std::vectorint64_t queryRopeShape {1, 2, 1, 64}; std::vectorint64_t keyRopeShape {1, 2, 1, 64}; std::vectorint64_t sinksShape {1}; int64_t queryShapeSize GetShapeSize(queryShape); int64_t keyShapeSize GetShapeSize(keyShape); int64_t valueShapeSize GetShapeSize(valueShape); int64_t sparseIndicesShapeSize GetShapeSize(sparseIndicesShape); int64_t attentionOutShapeSize GetShapeSize(attentionOutShape); int64_t softmaxMaxShapeSize GetShapeSize(softmaxMaxShape); int64_t softmaxSumShapeSize GetShapeSize(softmaxSumShape); int64_t queryRopeShapeSize GetShapeSize(queryRopeShape); int64_t keyRopeShapeSize GetShapeSize(keyRopeShape); int64_t sinksShapeSize GetShapeSize(sinksShape); std::vectorfloat queryHostData(queryShapeSize, 1); std::vectorfloat keyHostData(keyShapeSize, 1); std::vectorfloat valueHostData(valueShapeSize, 1); std::vectorint32_t sparseIndicesHostData(sparseIndicesShapeSize, 1); std::vectorfloat attentionOutHostData(attentionOutShapeSize, 1); std::vectorfloat softmaxMaxHostData(softmaxMaxShapeSize, 1); std::vectorfloat softmaxSumHostData(softmaxSumShapeSize, 1); std::vectorfloat queryRopeHostData(queryRopeShapeSize, 1); std::vectorfloat keyRopeHostData(keyRopeShapeSize, 1); std::vectorfloat sinksHostData(sinksShapeSize, 1); // Create query aclTensor. int ret CreateAclTensor(queryHostData, queryShape, resources.queryDeviceAddr, aclDataType::ACL_FLOAT16, resources.queryTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create key aclTensor. ret CreateAclTensor(keyHostData, keyShape, resources.keyDeviceAddr, aclDataType::ACL_FLOAT16, resources.keyTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create value aclTensor. ret CreateAclTensor(valueHostData, valueShape, resources.valueDeviceAddr, aclDataType::ACL_FLOAT16, resources.valueTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create sparseIndices aclTensor. ret CreateAclTensor(sparseIndicesHostData, sparseIndicesShape, resources.sparseIndicesDeviceAddr, aclDataType::ACL_INT32, resources.sparseIndicesTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create queryRope aclTensor. ret CreateAclTensor(queryRopeHostData, queryRopeShape, resources.queryRopeDeviceAddr, aclDataType::ACL_FLOAT16, resources.queryRopeTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create keyRope aclTensor. ret CreateAclTensor(keyRopeHostData, keyRopeShape, resources.keyRopeDeviceAddr, aclDataType::ACL_FLOAT16, resources.keyRopeTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create sinks aclTensor. ret CreateAclTensor(sinksHostData, sinksShape, resources.sinksDeviceAddr, aclDataType::ACL_FLOAT, resources.sinksTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create attention_out aclTensor. ret CreateAclTensor(attentionOutHostData, attentionOutShape, resources.attentionOutDeviceAddr, aclDataType::ACL_FLOAT16, resources.attentionOutTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create softmax_max aclTensor. ret CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, resources.softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, resources.softmaxMaxTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } // Create softmax_sum aclTensor. ret CreateAclTensor(softmaxSumHostData, softmaxSumShape, resources.softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, resources.softmaxSumTensor); if (!CHECK_RET(ret ACL_SUCCESS)) { return ret; } return ACL_SUCCESS; } int ExecuteSparseFlashAttention(TensorResources resources, aclrtStream stream, void** workspaceAddr, uint64_t* workspaceSize) { int64_t d 2; double scaleValue 1 / sqrt(d); int64_t sparseBlockSize 64; constexpr const char layerOutStr[] BSND; constexpr size_t layerOutLen sizeof(layerOutStr); char layoutQuery[layerOutLen]; char layoutKv[layerOutLen]; errno_t memcpyRet memcpy_s(layoutQuery, sizeof(layoutQuery), layerOutStr, layerOutLen); if (memcpyRet ! 0) { LOG_PRINT(memcpy_s layoutQuery failed. ERROR: %d\n, memcpyRet); return -1; } memcpyRet memcpy_s(layoutKv, sizeof(layoutKv), layerOutStr, layerOutLen); if (memcpyRet ! 0) { LOG_PRINT(memcpy_s layoutKv failed. ERROR: %d\n, memcpyRet); return -1; } int64_t sparseMode 3; int64_t preTokens 9223372036854775807; int64_t nextTokens 9223372036854775807; int64_t attentionMode 2; bool returnSoftmaxLse false; aclOpExecutor* executor; int ret aclnnSparseFlashAttentionV2GetWorkspaceSize(resources.queryTensor, resources.keyTensor, resources.valueTensor, resources.sparseIndicesTensor, nullptr, nullptr, nullptr, resources.queryRopeTensor, resources.keyRopeTensor, resources.sinks, scaleValue, sparseBlockSize, layoutQuery, layoutKv, sparseMode, preTokens, nextTokens, attentionMode, returnSoftmaxLse, resources.attentionOutTensor, resources.softmaxMaxTensor, resources.softmaxSumTensor, workspaceSize, executor); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclnnSparseFlashAttentionV2GetWorkspaceSize failed. ERROR: %d\n, ret); return ret; } if (*workspaceSize 0ULL) { ret aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret; } } ret aclnnSparseFlashAttentionV2(*workspaceAddr, *workspaceSize, executor, stream); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclnnSparseFlashAttentionV2 failed. ERROR: %d\n, ret); return ret; } return ACL_SUCCESS; } int PrintOutResult(std::vectorint64_t shape, void** deviceAddr) { auto size GetShapeSize(shape); std::vectoraclFloat16 resultData(size, 0); auto ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret; } for (int64_t i 0; i size; i) { LOG_PRINT(mean result[%ld] is: %f\n, i, aclFloat16ToFloat(resultData[i])); } return ACL_SUCCESS; } void CleanupResources(TensorResources resources, void* workspaceAddr, aclrtStream stream, int32_t deviceId) { if (resources.queryTensor) { aclDestroyTensor(resources.queryTensor); } if (resources.keyTensor) { aclDestroyTensor(resources.keyTensor); } if (resources.valueTensor) { aclDestroyTensor(resources.valueTensor); } if (resources.sparseIndicesTensor) { aclDestroyTensor(resources.sparseIndicesTensor); } if (resources.attentionOutTensor) { aclDestroyTensor(resources.attentionOutTensor); } if (resources.softmaxMaxTensor) { aclDestroyTensor(resources.softmaxMaxTensor); } if (resources.softmaxSumTensor) { aclDestroyTensor(resources.softmaxSumTensor); } if (resources.queryRopeTensor) { aclDestroyTensor(resources.queryRopeTensor); } if (resources.keyRopeTensor) { aclDestroyTensor(resources.keyRopeTensor); } if (resources.sinksTensor) { aclDestroyTensor(resources.sinksTensor); } if (resources.queryDeviceAddr) { aclrtFree(resources.queryDeviceAddr); } if (resources.keyDeviceAddr) { aclrtFree(resources.keyDeviceAddr); } if (resources.valueDeviceAddr) { aclrtFree(resources.valueDeviceAddr); } if (resources.sparseIndicesDeviceAddr) { aclrtFree(resources.sparseIndicesDeviceAddr); } if (resources.attentionOutDeviceAddr) { aclrtFree(resources.attentionOutDeviceAddr); } if (resources.softmaxMaxDeviceAddr) { aclrtFree(resources.softmaxMaxDeviceAddr); } if (resources.softmaxSumDeviceAddr) { aclrtFree(resources.softmaxSumDeviceAddr); } if (resources.queryRopeDeviceAddr) { aclrtFree(resources.queryRopeDeviceAddr); } if (resources.keyRopeDeviceAddr) { aclrtFree(resources.keyRopeDeviceAddr); } if (resources.sinksDeviceAddr) { aclrtFree(resources.sinksDeviceAddr); } if (workspaceAddr) { aclrtFree(workspaceAddr); } if (stream) { aclrtDestroyStream(stream); } aclrtResetDevice(deviceId); aclFinalize(); } } // namespace int main() { int32_t deviceId 0; aclrtStream stream nullptr; TensorResources resources {}; void* workspaceAddr nullptr; uint64_t workspaceSize 0; std::vectorint64_t attentionOutShape {1, 2, 1, 16}; std::vectorint64_t softmaxMaxShape {1, 2, 1, 16}; std::vectorint64_t softmaxSumShape {1, 2, 1, 16}; int ret ACL_SUCCESS; // 1. Initialize device and stream ret Init(deviceId, stream); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret; } // 2. Initialize tensors ret InitializeTensors(resources); if (!CHECK_RET(ret ACL_SUCCESS)) { CleanupResources(resources, workspaceAddr, stream, deviceId); return ret; } // 3. Execute the operation ret ExecuteSparseFlashAttention(resources, stream, workspaceAddr, workspaceSize); if (!CHECK_RET(ret ACL_SUCCESS)) { CleanupResources(resources, workspaceAddr, stream, deviceId); return ret; } // 4. Synchronize stream ret aclrtSynchronizeStream(stream); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); CleanupResources(resources, workspaceAddr, stream, deviceId); return ret; } // 5. Process results printf(-----------attentionOut输出-----------\n); PrintOutResult(attentionOutShape, resources.attentionOutDeviceAddr); printf(-----------softmaxMax输出-----------\n); PrintOutResult(softmaxMaxShape, resources.softmaxMaxDeviceAddr); printf(-----------softmaxSum输出-----------\n); PrintOutResult(softmaxSumShape, resources.softmaxSumDeviceAddr); // 6. Cleanup resources CleanupResources(resources, workspaceAddr, stream, deviceId); return 0; }【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考