BSASelectBlockMask【免费下载链接】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]Sq是query最大序列长度Skv是key最大序列长度, 则压缩后块数$$ 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参数说明参数名输入/输出/属性描述数据类型数据格式query输入注意力计算中的query矩阵即公式中的query。FLOAT16、BFLOAT16NDkey输入注意力计算中的key矩阵即公式中的key。数据类型与query保持一致NDblock_shape输入稀疏块形状数组指定每个稀疏块的二维尺寸行数和列数即公式中的blockShape。INT64-post_block_shape输入预留参数用于Softmax后二次压缩。--actual_seq_lengths输入每个batch的query的实际序列长度。INT64-actual_seq_lengths_kv输入key的实际序列长度。INT64-actual_block_len_query输入每个query block内实际压缩的有效seq长度用于部分压缩场景, 即公式中的actualBlockLenQuery。INT64-actual_block_len_key输入每个key block内实际压缩的有效seq长度用于部分压缩场景, 即公式中的actualBlockLenKey。INT64-q_input_layout输入query的数据排布格式。指示输入张量在内存中的具体排布。String-kv_input_layout输入key的数据排布格式。指示输入张量在内存中的具体排布。String-num_key_value_heads输入key的注意力头数。Int-scale_value输入缩放系数即公式中的scale。Float-sparsity输入稀疏度保留比例。指定公式中attn_score中需要保留的块位置占全部块位置的比例。取值范围(0.0, 1.0)。Float-block_sparse_mask_out输出块状稀疏掩码输出即公式中的blockSparseMaskOut。INT8ND约束说明actual_seq_lengths在queryLayout为TND时必选actual_seq_lengths_kv在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。block_sparse_mask_out数据类型为INT8二值0或1。post_block_shape当前不支持必须传入nullptr。actual_block_len_query / actual_block_len_key若非null每个元素取值范围[0, blockShapeX] / [0, blockShapeY]为null时完整压缩。调用说明调用方式调用样例说明aclnn APItest_aclnn_bsa_select_block_mask只支持MHA和格式为TND\BNSD的场景通过aclnnBSASelectBlockMask接口方式调用BSASelectBlockMask算子。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考