aclnnBSASelectBlockMask【免费下载链接】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训练系列产品×功能说明接口功能​aclnnBSASelectBlockMask是BSABlockSparseAttention的前置算子负责根据Query和Key的内容动态生成blockSparseMask使BSA的调用链从手动提供掩码变为根据Q/K内容自适应选择稀疏模式。计算公式​设blockShape [blockShapeX, blockShapeY]压缩后块数$$ Xblocks \lceil Sq / blockShapeX \rceil,\quad Yblocks \lceil Skv / blockShapeY \rceil $$Step1均值池化压缩 (Mean Pooling Compression)当actualBlockLenQuery / actualBlockLenKey为null时完整压缩$$ q_compressed[b, n, x, d] \frac{1}{blockShapeX} \sum_{i0}^{blockShapeX-1} query[b, n, x \cdot blockShapeX i, d] $$$$ k_compressed[b, n, y, d] \frac{1}{blockShapeY} \sum_{j0}^{blockShapeY-1} key[b, n, y \cdot blockShapeY j, d] $$当actualBlockLenQuery / actualBlockLenKey非null时部分压缩仅对每个block内前actualBlockLen个token取均值$$ q_compressed[b, n, x, d] \frac{1}{actualBlockLenQ[b,x]} \sum_{i0}^{actualBlockLenQ[b,x]-1} query[b, n, x \cdot blockShapeX i, d] $$$$ k_compressed[b, n, y, d] \frac{1}{actualBlockLenK[b,y]} \sum_{j0}^{actualBlockLenK[b,y]-1} key[b, n, y \cdot blockShapeY j, d] $$Step2aQK Matmul$$ score[b, n, x, y] scale \cdot \sum_{d0}^{D-1} q_compressed[b, n, x, d] \cdot k_compressed[b, n, y, d] $$Step2bSoftmax$$ attn_score[b, n, x, y] softmax(score[b, n, x, :]) \frac{\exp(score[b, n, x, y] - m_{final})}{l_{final}} $$Step3TopK选择生成索引$$ topk_value \text{round}(sparsity \times Xblocks \times Yblocks) $$$$ \mathcal{indices} \text{TopK}\left(attn_score[b, n, x, y],; topK_value\right) $$其中indices为attn_score[b, n, x, y] 中topk_value个最大值对应的索引集合。Step4生成BlockSparseMask$$ blockSparseMaskOut[b, n, x, y] \begin{cases} 1 (b, n, x, y) \in \mathcal{indices} \ 0 (b, n, x, y) \notin \mathcal{indices} \end{cases} $$BSASelectBlockMask输入query、key的数据排布格式支持从多种维度排布解读可通过queryLayout和keyLayout传入。为了方便理解后续支持的具体排布格式如BNSD、TND等此处先对排布格式中各缩写字母所代表的维度含义进行统一说明B表示输入样本批量大小BatchTB和S合轴紧密排列的长度Total tokensS表示输入样本序列长度Seq-LengthH表示隐藏层的大小Head-SizeN表示多头数Head-NumD表示隐藏层最小的单元尺寸需满足D H / NHead-Dim当前支持的布局queryLayout: TND BNSDkeyLayout: TND BNSD函数原型每个算子分为两段式接口必须先调用aclnnBSASelectBlockMaskGetWorkspaceSize接口获取计算所需workspace大小以及包含了算子计算流程的执行器再调用aclnnBSASelectBlockMask接口执行计算。aclnnStatus aclnnBSASelectBlockMaskGetWorkspaceSize( const aclTensor *query, const aclTensor *key, const aclIntArray *blockShape, const aclIntArray *postBlockShape, const aclIntArray *actualSeqLengths, const aclIntArray *actualSeqLengthsK, const aclIntArray *actualBlockLenQuery, const aclIntArray *actualBlockLenKey, char *queryLayout, char *keyLayout, int64_t numKeyHeads, double scaleValue, double sparsity, aclTensor *blockSparseMaskOut, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnBSASelectBlockMask( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)aclnnBSASelectBlockMaskGetWorkspaceSize参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorqueryaclTensor*输入注意力计算中的query矩阵即公式中的query。不支持空Tensor。支持的shape为TND: [totalQTokens, headNum, headDim]。BNSD: [batch, headNum, maxQSeqLength, headDim]。FLOAT16、BFLOAT16ND3-4×keyaclTensor*输入注意力计算中的key矩阵即公式中的key。不支持空Tensor。支持的shape为TND: [totalKTokens, numKeyHeads, headDim]。BNSD: [batch, numKeyHeads, maxKSeqLength, headDim]。FLOAT16、BFLOAT16ND3-4×blockShapeaclIntArray*输入稀疏块形状数组。指定每个稀疏块的二维尺寸行数和列数。当配置此输入时的元素要求必须包含至少两个元素 [blockShapeX, blockShapeY]。blockShapeX: Q方向块大小必须为64的倍数且大于0。blockShapeY: KV方向块大小必须为64的倍数且大于0。如不配置传nullptr算子将默认blockShapeX 128blockShapeY 128。INT64-1-postBlockShapeaclIntArray*输入预留参数用于Softmax后二次压缩。当前不支持必须传入nullptr。----actualSeqLengthsaclIntArray*输入每个batch的query的实际序列长度。用于描述变长序列场景下即含有Padding填充数据的场景每个Batch中实际有效的query token数量。变长序列场景当queryLayout为 TND 时该项输入必须配置。定长/变长场景当queryLayout为 BNSD 时如配置该项算子会按指定的有效长度处理忽略Padding部分的数据提升性能如不配置传nullptr算子将默认把query shape中的S维度作为有效长度进行全量处理。INT64-1-actualSeqLengthsKaclIntArray*输入key的实际序列长度数组。用于描述变长序列场景下即含有Padding填充数据的场景每个Batch中实际有效的key token数量。变长序列场景当keyLayout为 TND 时该项输入必须配置。定长/变长场景当keyLayout为 BNSD 时如配置该项算子会按指定的有效长度处理忽略Padding部分的数据提升性能如不配置传nullptr算子将默认把key shape中的S维度作为有效长度进行全量处理。INT64-1-actualBlockLenQueryaclIntArray*输入每个query block内实际压缩的有效seq长度。用于部分压缩场景如末尾不完整块或仅压缩有效token。可选输入BNSD场景shape为 [B, Xblocks]。TND场景shape为 [TotalBlockNum_Q]各batch的Xblocks堆叠。每个元素取值范围 [0, blockShapeX]。当actualBlockLen 0时对应block的q_cmp填0向量不会被topK选中。当actualBlockLen 0时仅对前actualBlockLen个token取均值。如不配置传nullptr对query进行完整压缩使用完整blockShapeX长度。INT64-1-actualBlockLenKeyaclIntArray*输入每个key block内实际压缩的有效seq长度。用于部分压缩场景如末尾不完整块或仅压缩有效token。可选输入BNSD场景shape为 [B, Yblocks]。TND场景shape为 [TotalBlockNum_KV]各batch的Yblocks堆叠。每个元素取值范围 [0, blockShapeY]。当actualBlockLen 0时对应block的k_cmp填0向量 不会被topK选中。当actualBlockLen 0时仅对前actualBlockLen个token取均值。如不配置传nullptr对key进行完整压缩使用完整blockShapeY长度。INT64-1-queryLayoutchar*输入query的数据排布格式。指示输入张量在内存中的具体排布。当前仅支持 TND、BNSDqueryLayout与keyLayout需要保持一致。----keyLayoutchar*输入key的数据排布格式。指示输入张量在内存中的具体排布。当前仅支持 TND、BNSDqueryLayout与keyLayout需要保持一致。----numKeyHeadsint64_t输入key的注意力头数。当前仅支持MHAnumKeyHeads必须等于numHeads。----scaleValuedouble输入缩放系数即公式中的scale。用于注意力分数的归一化处理。一般设置为1 / sqrt(D)。----sparsitydouble输入稀疏度保留比例。指定公式中attn_score中需要保留的块位置占全部块位置的比例。取值范围 (0.0, 1.0)。----blockSparseMaskOutaclTensor*输出块状稀疏掩码输出表示根据Q/K内容自适应生成的稀疏pattern。可直接作为BSA算子的blockSparseMask输入。不支持空Tensor。shape为 [B, N, Xblocks, Yblocks]Xblocks ceilDiv(maxQSeqlen, blockShapeX)。Yblocks ceilDiv(maxKSeqlen, blockShapeY)。值为1表示该block参与注意力计算值为0表示不参与。INT8ND4×workspaceSizeuint64_t*输出返回需要在Device侧申请的workspace大小。-----executoraclOpExecutor**输出返回op执行器包含了算子计算流程。-----返回值aclnnStatus返回状态码具体参见aclnn返回码。第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入query传入的是空指针。输入key传入的是空指针。blockSparseMaskOut传入的是空指针。ACLNN_ERR_PARAM_INVALID161002query、key或者blockSparseMaskOut数据类型不在支持的范围之内。query和key数据类型不一致。aclnnBSASelectBlockMask参数说明参数名输入/输出描述workspace输入在Device侧申请的workspace内存地址。workspaceSize输入在Device侧申请的workspace大小由第一段接口aclnnBSASelectBlockMaskGetWorkspaceSize获取。executor输入op执行器包含了算子计算流程。stream输入指定执行任务的AscendCL stream流。返回值aclnnStatus返回状态码具体参见aclnn返回码。约束说明该接口若与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。actualSeqLengths在queryLayout为 TND 时必选actualSeqLengthsK在keyLayout为 TND 时必选。根据算子支持的输入Layoutquery张量Shape中对应的head维度大小记为N1key张量Shape中对应的head维度大小记为N2。必须满足N1 N2仅支持MHA。headDim 128。blockShapeX和blockShapeY必须为64的倍数。query和key压缩后query和key对应的Xblocks和Yblocks需满足Xblocks * Yblocks 1。query和key的数据类型必须一致仅支持FLOAT16和BFLOAT16。blockSparseMaskOut数据类型为INT8二值0或1。postBlockShape当前不支持必须传入nullptr。actualBlockLenQuery / actualBlockLenKey若非null每个元素取值范围 [0, blockShapeX] / [0, blockShapeY]为null时完整压缩。调用示例示例代码如下仅供参考具体编译和执行过程请参考编译与运行样例。#include iostream #include vector #include cstring #include cmath #include cstdint #include acl/acl.h #include aclnn/opdev/fp16_t.h #include aclnnop/aclnn_bsa_select_block_mask.h using namespace std; #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; } int Init(int32_t deviceId, 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 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) { if (shape.empty()) { LOG_PRINT(CreateAclTensor: ERROR - shape is empty\n); return -1; } for (size_t i 0; i shape.size(); i) { if (shape[i] 0) { LOG_PRINT(CreateAclTensor: ERROR - shape[%zu]%ld is invalid\n, i, shape[i]); return -1; } } auto size GetShapeSize(shape) * sizeof(T); if (hostData.size() ! static_castsize_t(GetShapeSize(shape))) { LOG_PRINT(CreateAclTensor: ERROR - hostData size mismatch: %zu vs %ld\n, hostData.size(), GetShapeSize(shape)); return -1; } *deviceAddr nullptr; 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); 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); aclrtFree(*deviceAddr); *deviceAddr nullptr; return ret); std::vectorint64_t strides(shape.size(), 1); if (shape.size() 1) { for (int64_t i static_castint64_t(shape.size()) - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } } *tensor nullptr; *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); CHECK_RET(*tensor ! nullptr, LOG_PRINT(aclCreateTensor failed - returned nullptr\n); aclrtFree(*deviceAddr); *deviceAddr nullptr; return -1); return 0; } int main() { // 1. device/stream初始化 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. 设置核心参数 (以BNSD Layout为例) int32_t batch 1; int32_t numHeads 4; int32_t numKHeads 4; // 仅支持MHA: N KN int32_t qSeqlen 256 * 1024; // 256K int32_t kSeqlen 256 * 1024; // 256K int32_t headDim 128; int32_t blockShapeX 128; int32_t blockShapeY 128; // 块数量计算 int32_t ceilQ (qSeqlen blockShapeX - 1) / blockShapeX; int32_t ceilK (kSeqlen blockShapeY - 1) / blockShapeY; // 3. 构建张量Shape std::vectorint64_t qShape {batch, numHeads, qSeqlen, headDim}; std::vectorint64_t kShape {batch, numKHeads, kSeqlen, headDim}; std::vectorint64_t maskShape {batch, numHeads, ceilQ, ceilK}; // 4. 分配并初始化Host数据 int64_t qSize GetShapeSize(qShape); int64_t kSize GetShapeSize(kShape); std::vectorop::fp16_t qData(qSize, 0.1f); std::vectorop::fp16_t kData(kSize, 0.1f); // mask输出初始化为0 std::vectorint8_t maskOutData(GetShapeSize(maskShape), 0); // 创建所有aclTensor void *qAddr nullptr, *kAddr nullptr, *maskAddr nullptr; aclTensor *qTensor nullptr, *kTensor nullptr, *maskTensor nullptr; CreateAclTensor(qData, qShape, qAddr, aclDataType::ACL_FLOAT16, qTensor); CreateAclTensor(kData, kShape, kAddr, aclDataType::ACL_FLOAT16, kTensor); CreateAclTensor(maskOutData, maskShape, maskAddr, aclDataType::ACL_INT8, maskTensor); // 5. 创建aclIntArray属性参数 std::vectorint64_t blockShapeVec {blockShapeX, blockShapeY}; aclIntArray *blockShapeArr aclCreateIntArray(blockShapeVec.data(), blockShapeVec.size()); // 6. 标量与字符串参数配置 char queryLayoutBuffer[16] BNSD; char keyLayoutBuffer[16] BNSD; double scaleValue 1.0 / std::sqrt(static_castdouble(headDim)); double sparsity 0.5; // 7. 调用第一段接口: GetWorkspaceSize uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; LOG_PRINT(Calling aclnnBSASelectBlockMaskGetWorkspaceSize...\n); ret aclnnBSASelectBlockMaskGetWorkspaceSize( qTensor, kTensor, blockShapeArr, nullptr, nullptr, nullptr, nullptr, nullptr, queryLayoutBuffer, keyLayoutBuffer, static_castint64_t(numKHeads), scaleValue, sparsity, maskTensor, workspaceSize, executor ); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(GetWorkspaceSize failed. ERROR: %d\n, ret); return ret); CHECK_RET(executor ! nullptr, LOG_PRINT(executor is null after GetWorkspaceSize\n); return -1); LOG_PRINT(Workspace size required: %lu bytes\n, workspaceSize); // 8. 分配workspace 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); } // 9. 调用第二段接口: 执行计算 LOG_PRINT(Calling aclnnBSASelectBlockMask...\n); ret aclnnBSASelectBlockMask(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnBSASelectBlockMask failed. ERROR: %d\n, ret); return ret); // 10. 同步Stream等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 11. 将结果拷贝回Host侧打印 (blockSparseMaskOut) int64_t maskSize GetShapeSize(maskShape); ret aclrtMemcpy(maskOutData.data(), maskSize * sizeof(int8_t), maskAddr, maskSize * sizeof(int8_t), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed.\n); return ret); LOG_PRINT(Execution Success! BlockSparseMask output (first 20 elements):\n); int64_t displayCount (maskSize 20) ? maskSize : 20; for (int64_t i 0; i displayCount; i) { LOG_PRINT( mask index %ld: %u\n, i, static_castunsigned int(maskOutData[i])); } // 12. 释放所有资源 LOG_PRINT(Cleaning up resources...\n); if (workspaceAddr) { aclrtFree(workspaceAddr); } aclrtFree(qAddr); aclrtFree(kAddr); aclrtFree(maskAddr); aclDestroyTensor(qTensor); aclDestroyTensor(kTensor); aclDestroyTensor(maskTensor); aclDestroyIntArray(blockShapeArr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); LOG_PRINT(BSASelectBlockMask Test completed successfully!\n); return 0; }【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考