当前位置：首页 > news >正文

CANN块稀疏注意力算子

news 2026/5/9 13:19:00

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两种精度
非对齐序列: 支持序列长度不能被块大小整除的场景,自动处理边界情况

接口定义

aclnnBlockSparseAttentionGetWorkspaceSize

aclnnStatus 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); // 返回executor

aclnnBlockSparseAttention

aclnnStatus 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稀疏pattern

BlockSparseAttention使用blockSparseMask稀疏pattern,不需要selectIdx索引格式和传统的sBlockIdx计算:

blockSparseMask: 形状为[batch, numHeads, qBlockNum, kvBlockNum]的稀疏Mask
blockShapeX: Q方向基本块大小
blockShapeY: KV方向基本块大小

示例

对于qSeqlen=512, kvSeqlen=1024, blockShapeX=128, blockShapeY=128:

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]，且值必须大于0，blockShapeY必须为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”时必选。
当前不支持PagedAttention，blockTableOptional当前只支持传入nullptr。
query输入为BFLOAT16时，innerPrecise只能配置为0。