CANN块稀疏注意力算子

BlockSparseAttention算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer概述BlockSparseAttention是一个基于CATLASS模板库实现的高性能稀疏注意力算子支持灵活的块级稀疏模式。功能特性灵活的稀疏块模式: 支持自定义的x*y块级稀疏模式,通过blockShape参数指定多种KV Cache布局: 支持TND [T, N, D]和BNSD [B, N, S, D]的的KV cache布局多种Q输入布局: 支持TND [T, N, D]和BNSD [B, N, S, D]两种Query输入布局高性能计算: 基于CATLASS模板库,充分利用昇腾A2硬件特性多精度支持: 支持FP16和BF16两种精度非对齐序列: 支持序列长度不能被块大小整除的场景,自动处理边界情况接口定义aclnnBlockSparseAttentionGetWorkspaceSizeaclnnStatus aclnnBlockSparseAttentionGetWorkspaceSize( const aclTensor *query, // Query输入 const aclTensor *key, // Key输入 const aclTensor *value, // Value输入 const aclTensor *blockSparseMaskOptional, // 稀疏Mask const aclTensor *attenMaskOptional, // Attention mask (当前不支持) const aclIntArray *blockShapeOptional, // 稀疏块形状数组 const aclIntArray *actualSeqLengthsOptional, // 实际Q序列长度 (TND格式时必选) const aclIntArray *actualSeqLengthsKvOptional, // 实际KV序列长度 (TND格式时必选) const aclTensor *blockTableOptional, // Block表 (用于PagedAttention当前不支持) char *qInputLayout, // Query的数据排布格式 (TND/BNSD) char *kvInputLayout, // Key/Value的数据排布格式 (TND/BNSD) int64_t numKeyValueHeads, // KV头数 int64_t maskType, // Mask类型 double scaleValue, // 缩放因子 int64_t innerPrecise, // Softmax计算采取的精度级别 int64_t blockSize, // Block大小 (用于PagedAttention当前不支持) int64_t preTokens, // 滑窗参数 (当前不支持) int64_t nextTokens, // 滑窗参数 (当前不支持) int64_t softmaxLseFlag, // 是否输出LSE aclTensor *attentionOut, // 输出tensor aclTensor *softmaxLseOptional, // Softmax LSE输出 (可选) uint64_t *workspaceSize, // 返回workspace大小 aclOpExecutor **executor); // 返回executoraclnnBlockSparseAttentionaclnnStatus aclnnBlockSparseAttention( void *workspace, // workspace地址 uint64_t workspaceSize, // workspace大小 aclOpExecutor *executor, // executor aclrtStream stream); // ACL stream参数说明输入参数query: Query tensor,支持以下布局TND: [total_tokens, num_heads, head_dim]BNSD: [batch, num_heads, seqlen, head_dim]key: Key tensor,支持以下布局TND: [total_kv_tokens, num_kv_heads, head_dim]BNSD: [batch, num_kv_heads, kv_seqlen, head_dim]value: Value tensor,支持以下布局TND: [total_kv_tokens, num_kv_heads, head_dim]BNSD: [batch, num_kv_heads, kv_seqlen, head_dim]blockSparseMaskOptional:稀疏Mask: [batch, headNum, ceilDiv(maxQSeqLength, blockShapeX), ceilDiv(maxKvSeqLength, blockShapeY)]blockShapeOptional:blockShapeX: Q方向块大小blockShapeY: KV方向块大小输出参数attentionOut: 注意力输出,shape与query相同softmaxLse: Softmax log-sum-exp输出 (可选)使用示例// 1. 准备输入数据 int64_t batch 1; int64_t qSeqlen 256; int64_t kvSeqlen 512; int64_t numHeads 8; int64_t numKvHeads 8; int64_t headDim 128; int64_t blockShapeX 128; int64_t blockShapeY 128; // 2. 创建tensor (略) // 3. 设置blockShape int64_t blockShapeData[2] {blockShapeX, blockShapeY}; aclIntArray *blockShape aclCreateIntArray(blockShapeData, 2); // 4. 调用算子 uint64_t workspaceSize 0; aclOpExecutor *executor nullptr; aclnnStatus ret aclnnBlockSparseAttentionGetWorkspaceSize( query, key, value, blockSparseMask, nullptr, blockShape, nullptr, nullptr, nullptr, BNSD, BNSD, numKvHeads, 0, scaleValue, 0, 128, 2147483647, 2147483647, 0, attentionOut, nullptr, workspaceSize, executor); if (ret ACLNN_SUCCESS) { void *workspace nullptr; aclrtMalloc(workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); ret aclnnBlockSparseAttention(workspace, workspaceSize, executor, stream); aclrtFree(workspace); }稀疏模式说明BlockSparseMask稀疏patternBlockSparseAttention使用blockSparseMask稀疏pattern,不需要selectIdx索引格式和传统的sBlockIdx计算:blockSparseMask: 形状为[batch, numHeads, qBlockNum, kvBlockNum]的稀疏MaskblockShapeX: Q方向基本块大小blockShapeY: KV方向基本块大小示例对于qSeqlen512, kvSeqlen1024, blockShapeX128, blockShapeY128:Q方向: ceil(512/128)4块 [0, 1, 2, 3] KV方向: ceil(1024/128)8块 [0, 1, 2, 3, 4, 5, 6, 7] 每个Q块可以选择任意KV块组合 - Q块0选择的KV块: [0, 2, 5] - Q块1选择的KV块: [1, 3, 6] - Q块2选择的KV块: [0, 4, 7] - Q块3选择的KV块: [2, 5, 6]性能特点内存效率: 相比密集注意力,显著减少内存使用计算效率: 只计算指定的稀疏块,减少计算量灵活配置: 支持任意blockShape的稀疏块配置硬件优化: 充分利用昇腾A2的AI Core和Vector Core注意事项qInputLayout当前仅支持TND和BNSD。kvInputLayout当前仅支持TND和BNSD。blockShapeOptional如果传入则必须包含至少两个元素[blockShapeX, blockShapeY]且值必须大于0blockShapeY必须为128的倍数。qSeqlen和kvSeqlen不需要被blockShape整除支持非对齐场景实际分块数通过向上取整计算。blockSparseMaskOptional当前必须传入且shape必须为[batch, headNum, ceilDiv(maxQS, blockShapeX), ceilDiv(maxKVS, blockShapeY)]。attentionMaskOptional当前只支持传入nullptr。maskType当前只支持输入0表示不加mask。actualSeqLengthsOptional在qInputLayout为“TND”时必选actualSeqLengthsKvOptional在kvInputLayout为“TND”时必选。当前不支持PagedAttentionblockTableOptional当前只支持传入nullptr。query输入为BFLOAT16时innerPrecise只能配置为0。编译说明算子编译集成在ops-transformer的统一编译流程中:bash build.sh版本历史v1.0.0: 初始版本【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考