CANN/ops-transformer FlashAttention变长分数计算V5

aclnnFlashAttentionVarLenScoreV5【免费下载链接】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 训练系列产品×功能说明接口功能训练场景下使用FlashAttention算法实现self-attention自注意力的计算。对标竞品适配gptoss模型支持sink功能。跟aclnnFlashAttentionVarLenScoreV3接口的区别是增加sinkInOptional可选输入,保留了aclnnFlashAttentionVarLenScoreV4的softmaxOutLayout可选输入。Ascend 950PR/Ascend 950DT 暂不支持sinkInOptional与softmaxOutLayout参数。计算公式 注意力的正向计算公式如下$$ AttentionDropout(Softmax(Mask(scale*(querykey^T queryRopekeyRope^T) pse),atten_mask),keep_prob)*value $$其中增加sink之后计算逻辑见下主要修改相关softmax_max和softmax_sum逻辑计算部分$$ S Q * K^{T} $$$$ m max(sink, max(S)) $$$$ Attention \frac{e^{S - m} * V}{\sum e^{S-m} e^{sink - m}} $$函数原型每个算子分为两段式接口必须先调用“aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器再调用“aclnnFlashAttentionVarLenScoreV5”接口执行计算。aclnnStatus aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize( const aclTensor *query, const aclTensor *queryRope, const aclTensor *key, const aclTensor *keyRope, const aclTensor *value, const aclTensor *realShiftOptional, const aclTensor *dropMaskOptional, const aclTensor *paddingMaskOptional, const aclTensor *attenMaskOptional, const aclTensor *sinkOptional, 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, char *softmaxOutLayout, const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, const aclTensor *softmaxOutOut, const aclTensor *attentionOutOut, uint64_t *workspaceSize, aclOpExecutor **executor);aclnnStatus aclnnFlashAttentionVarLenScoreV5( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorquery输入公式中的query。数据类型与key/value的数据类型一致。FLOAT16、BFLOAT16、FLOAT32ND[TND]√queryRope可选输入公式中的queryRope。数据类型与key/value的数据类型一致。BFLOAT16ND[TND]√key输入公式中的key。数据类型与query/value的数据类型一致。FLOAT16、BFLOAT16、FLOAT32ND[TND]√keyRope可选输入公式中的keyRope。数据类型与query/value的数据类型一致。BFLOAT16ND[TND]√value输入公式中的value。数据类型与query/key的数据类型一致。FLOAT16、BFLOAT16、FLOAT32ND[TND]√realShiftOptional可选输入公式中的pse。与queryRope、keyRope不兼容。不用rope时数据类型和query一致需要与pseType配套使用。FLOAT16、BFLOAT16、FLOAT32ND[B,N,1024,Skv]、[1,N,1024,Skv]、[pseTotalLen]、[B,N]、[N]√dropMaskOptional可选输入公式中的Dropout。与queryRope、keyRope不兼容。UINT8ND0、1√paddingMaskOptional可选输入预留参数暂未使用。-----attenMaskOptional可选输入公式中的atten_mask。取值为1代表该位不参与计算为0代表该位参与计算。BOOL、UINT8ND[B,N,Sq,Skv]、[B,1,Sq,Skv]、[1,1,Sq,Skv]、[Sq,Skv]√sinkOptional可选输入公式中的sink。提供sink功能。FLOAT32ND[headNum]√prefixOptional可选输入代表prefix稀疏计算场景每个Batch的N值。-INT64ND0、1-actualSeqQLenOptional可选输入描述了每个Batch对应的query的sequence length。-INT64ND1-actualSeqKvLenOptional可选输入描述了每个Batch对应的key/value的sequence length。-INT64ND1-qStartIdxOptional可选输入代表外切场景当前分块的query的sequence在全局中的起始索引。-INT64ND0、1-kvStartIdxOptional可选输入代表外切场景当前分块的key和value的sequence在全局中的起始索引。-INT64ND0、1-scaleValue可选输入公式中的scale代表缩放系数。-DOUBLE---keepProb可选输入代表dropMaskOptional中1的比例。取值范围为(0, 1]。DOUBLE---preTokens可选输入用于稀疏计算 表示sliding window的左边界。-INT64---nextTokens可选输入用于稀疏计算表示sliding window的右边界。-INT64---headNum输入代表单卡的head个数即输入query的N轴长度。-INT64---inputLayout输入代表输入query、key、value的数据排布格式。支持TND。String---innerPrecise可选输入用于提升精度。默认配置为0即可。INT64---sparseMode可选输入表示sparse的模式。支持配置值为支持配置0~8不支持5。传入rope时不支持6。INT64---pseType可选输入控制mul与add计算顺序。无rope时配置为0-3有rope仅支持配置值为1。INT64---softmaxOutLayout可选输入用于控制TND场景下softmax输出。传入same_as_input时softmax输出排布与输入保持一致为TND排布传入空字符串时与原逻辑保持一致softmax输出排布为NTD。String---softmaxMaxOut输出Softmax计算的Max中间结果用于反向计算。输出shape由softmaxOutLayout决定。FLOATND[N,T,8]或[T,N,8]√softmaxSumOut输出Softmax计算的Sum中间结果用于反向计算。输出shape由softmaxOutLayout决定。FLOATND[N,T,8]或[T,N,8]√softmaxOutOut输出预留参数暂未使用。-----attentionOutOut输出计算公式的最终输出。数据类型和shape类型与query保持一致。BFLOAT16ND[T,N,D]√workspaceSize输出返回需要在Device侧申请的workspace大小。-----executor输出返回op执行器包含了算子计算流程。-----返回值aclnnStatus返回状态码具体参见aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入参数是必选输入输出或者必选属性且是空指针。ACLNN_ERR_PARAM_INVALID161002query、key、value、realShiftOptional、dropMaskOptional、paddingMaskOptional、attenMaskOptional、sinkOptional、softmaxMaxOut、softmaxSumOut、softmaxOutOut、attentionOutOut的数据类型不在支持的范围内。query、key、value、realShiftOptional、dropMaskOptional、paddingMaskOptional、attenMaskOptional、sinkOptional、softmaxMaxOut、softmaxSumOut、softmaxOutOut、attentionOutOut的数据格式不在支持的范围内。ACLNN_ERR_INNER_NULLPTR561103API内部校验错误通常由于输入的shape或属性的规格不在支持的范围之内导致。aclnnFlashAttentionVarLenScoreV5参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的Stream。返回值返回aclnnStatus状态码具体参见aclnn返回码。约束说明确定性计算aclnnFlashAttentionVarLenScoreV5默认确定性实现。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。输入query、key、value的Bbatchsize必须相等。输入query、key、value的DHead-Dim必须满足qD kD kD vD。输入query、key、value的input_layout必须一致。输入query、key、value、realShiftOptional的数据类型必须一致。支持输入query的N和key/value的N不相等但必须成比例关系即Nq/Nkv必须是非0整数Nq取值范围1~256。当Nq/Nkv 1时即为GQAgrouped-query attention当Nkv1时即为MQAmulti-query attention。本文如无特殊说明N表示的是Nq。关于数据shape的约束其中T(B*S)取值范围为1~1M。B取值范围为1~20000。带prefixOptional的时候B最大支持1K。N取值范围为1~256。S取值范围为1~1M。D取值范围为1~768。query、key、value数据排布格式仅支持TNDT是B和S合轴紧密排列的数据每个batch的SeqLenQ和SeqLenKV其中BBatch表示输入样本批量大小、SSeq-Length表示输入样本序列长度、HHead-Size表示隐藏层的大小、NHead-Num表示多头数、DHead-Dim表示隐藏层最小的单元尺寸且满足DH/N。realShiftOptional如果Sq大于1024且每个batch的Sq与Skv等长且是sparseMode为0、2、3的下三角掩码场景可使能alibi位置编码压缩此时只需要输入原始PSE最后1024行实现内存优化即alibi_compress ori_pse[:, :, -1024:, :]具体如下参数每个batch不相同时shape为BNHSkv(H1024)。每个batch相同时shape为1NHSkv(H1024)。TND场景下每个batch段内部仍按[N, Sq_i, Skv_i]生成但存储与传参时统一flatten。若第i个batch段的真实query长度为Sq_i、真实key/value长度为Skv_i则该段PSE元素个数为N * Sq_i * Skv_i整段PSE总长度pseTotalLen为sum_i(N * Sq_i * Skv_i)。如果pseType为2或3的时候数据类型需为FLOAT32, 对应shape支持范围是[B,N]或[N]。如果不使能该参数realShiftOptional需要传入nullptrpseType需要传入1。innerPrecise当前0、1为保留配置值2为使能无效行计算其功能是避免在计算过程中存在整行mask进而导致精度有损失但是该配置会导致性能下降。 如果算子可判断出存在无效行场景会自动使能无效行计算例如sparseMode为3Sq Skv场景。sparseMode的约束如下:当所有的attenMaskOptional的shape小于2048且相同的时候建议使用default模式来减少内存使用量。配置为1、2、3、6时用户配置的preTokens、nextTokens不会生效。配置为0、4、7时须保证attenMaskOptional与preTokens、nextTokens的范围一致。用户不特意指定时建议传入0。sparse不同模式的详细说明请参见sparse模式说明。为1、2、3、4、6、7、8时应传入对应正确的attenMaskOptional否则将导致计算结果错误。当attenMaskOptional输入为None时sparseMode、preTokens、nextTokens参数不生效固定为全计算。配置为3时不支持无效行计算需要满足每个batch的SqSkv。配置为7时不支持可选输入realShiftOptional。配置为8时当每个sequence的q、kv等长时支持可选输入realShiftOptional针对全局做pse生成。支持q方向进行外切需要外切前每个sequence的q、kv等长外切后传入的actualSeqQLenOptional[0] - actualSeqKvLenOptional[0] qStartIdxOptional - kvStartIdxOptional 0本功能属实验性功能。部分场景下如果计算量过大可能会导致算子执行超时aicore error类型报错errorStr为timeout or trap error此时建议做轴切分处理注这里的计算量会受B、S、N、D等参数的影响值越大计算量越大。prefixOptional稀疏计算场景即sparseMode6当Sq Skv时prefix的N值取值范围[0, Skv]当Sq Skv时prefix的N值取值范围[Skv-Sq, Skv]。band场景preTokens和nextTokens之间必须要有交集。attenMaskOptional输入不支持补pad即attenMaskOptional中不能存在某一行全1的场景。actualSeqQLenOptional输入支持某个Batch上的S长度为0此时不支持可选输入realShiftOptional假设真实的S长度为[2,2,0,2,2]则传入的actualSeqQLenOptional为[2,4,4,6,8]。softmaxOutLayout支持传入空字符串、same_as_input。调用示例调用示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_flash_attention_score.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::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } void PrintOutResult(std::vectorint64_t shape, void** deviceAddr) { auto size GetShapeSize(shape); std::vectorfloat 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) { // 固定写法资源初始化 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::vectorT hostData, const std::vectorint64_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::vectorint64_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/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId 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. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t qShape {256, 1, 128}; std::vectorint64_t qRopeShape {256, 1, 64}; std::vectorint64_t kShape {256, 1, 128}; std::vectorint64_t kRopeShape {256, 1, 64}; std::vectorint64_t vShape {256, 1, 128}; std::vectorint64_t attenmaskShape {256, 256}; std::vectorint64_t sinkShape {1}; std::vectorint64_t attentionOutShape {256, 1, 128}; std::vectorint64_t softmaxMaxShape {256, 1, 8}; std::vectorint64_t softmaxSumShape {256, 1, 8}; void* qDeviceAddr nullptr; void* qRopeDeviceAddr nullptr; void* kDeviceAddr nullptr; void* kRopeDeviceAddr nullptr; void* vDeviceAddr nullptr; void* attenmaskDeviceAddr nullptr; void* sinkDeviceAddr nullptr; void* attentionOutDeviceAddr nullptr; void* softmaxMaxDeviceAddr nullptr; void* softmaxSumDeviceAddr nullptr; aclTensor* q nullptr; aclTensor* qRope nullptr; aclTensor* k nullptr; aclTensor* kRope nullptr; aclTensor* v nullptr; aclTensor* pse nullptr; aclTensor* dropMask nullptr; aclTensor* padding nullptr; aclTensor* attenmask nullptr; aclTensor* sink nullptr; aclTensor* attentionOut nullptr; aclTensor* softmaxMax nullptr; aclTensor* softmaxSum nullptr; aclTensor* softmaxOut nullptr; std::vectorfloat qHostData(32768, 1); std::vectorfloat qRopeHostData(16384, 1); std::vectorfloat kHostData(32768, 1); std::vectorfloat kRopeHostData(16384, 1); std::vectorfloat vHostData(32768, 1); std::vectoruint8_t attenmaskHostData(65536, 0); std::vectorfloat sinkHostData(1, 3.0); std::vectorfloat attentionOutHostData(32768, 0); std::vectorfloat softmaxMaxHostData(2048, 3.0); std::vectorfloat softmaxSumHostData(2048, 3.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(attenmaskHostData, attenmaskShape, attenmaskDeviceAddr, aclDataType::ACL_UINT8, attenmask); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(sinkHostData, sinkShape, sinkDeviceAddr, aclDataType::ACL_FLOAT, sink); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(attentionOutHostData, attentionOutShape, attentionOutDeviceAddr, aclDataType::ACL_BF16, attentionOut); 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); std::vectorint64_t prefixOp {0}; aclIntArray *prefix aclCreateIntArray(prefixOp.data(), 1); std::vectorint64_t qStartIdxOp {0}; std::vectorint64_t kvStartIdxOp {0}; aclIntArray *qStartIdx aclCreateIntArray(qStartIdxOp.data(), 1); aclIntArray *kvStartIdx aclCreateIntArray(kvStartIdxOp.data(), 1); std::vectorint64_t acSeqQLenOp {256}; std::vectorint64_t acSeqKvLenOp {256}; aclIntArray* acSeqQLen aclCreateIntArray(acSeqQLenOp.data(), acSeqQLenOp.size()); aclIntArray* acSeqKvLen aclCreateIntArray(acSeqKvLenOp.data(), acSeqKvLenOp.size()); 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}; char softmaxOutLayout[] same_as_input; // 3. 调用CANN算子库API需要修改为具体的Api名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnFlashAttentionVarLenScoreV5第一段接口 ret aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize( q, qRope, k, kRope, v, pse, dropMask, padding, attenmask, sink, prefix, acSeqQLen, acSeqKvLen, qStartIdx, kvStartIdx, scaleValue, keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise, sparseMode, pseType, softmaxOutLayout, softmaxMax, softmaxSum, softmaxOut, attentionOut, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnFlashAttentionVarLenScoreV5GetWorkspaceSize 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); } // 调用aclnnFlashAttentionVarLenScoreV5第二段接口 ret aclnnFlashAttentionVarLenScoreV5(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnFlashAttentionVarLenScoreV5 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(attentionOutShape, attentionOutDeviceAddr); PrintOutResult(softmaxMaxShape, softmaxMaxDeviceAddr); PrintOutResult(softmaxSumShape, softmaxSumDeviceAddr); // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(q); aclDestroyTensor(qRope); aclDestroyTensor(k); aclDestroyTensor(kRope); aclDestroyTensor(v); aclDestroyTensor(attenmask); aclDestroyTensor(sink); aclDestroyTensor(attentionOut); aclDestroyTensor(softmaxMax); aclDestroyTensor(softmaxSum); // 7. 释放device资源 aclrtFree(qDeviceAddr); aclrtFree(qRopeDeviceAddr); aclrtFree(kDeviceAddr); aclrtFree(kRopeDeviceAddr); aclrtFree(vDeviceAddr); aclrtFree(attenmaskDeviceAddr); aclrtFree(sinkDeviceAddr); aclrtFree(attentionOutDeviceAddr); aclrtFree(softmaxMaxDeviceAddr); aclrtFree(softmaxSumDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考