CANN/ops-transformer融合推理注意力分数算子

发布时间:2026/7/9 10:48:30

CANN/ops-transformer融合推理注意力分数算子 FusedInferAttentionScore【免费下载链接】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算子既可以支持全量计算场景PromptFlashAttention也可支持增量计算场景IncreFlashAttention。计算公式self-attention自注意力利用输入样本自身的关系构建了一种注意力模型。其原理是假设有一个长度为$n$的输入样本序列$x$$x$的每个元素都是一个$d$维向量可以将每个$d$维向量看作一个token embedding将这样一条序列经过3个权重矩阵变换得到3个维度为$n*d$的矩阵。self-attention的计算公式一般定义如下其中$Q$、$K$、$V$为输入样本的重要属性元素是输入样本经过空间变换得到且可以统一到一个特征空间中。公式及算子名称中的Attention为self-attention的简写。$$ Attention(Q,K,V)Score(Q,K)V $$本算子中Score函数采用Softmax函数self-attention计算公式为$$ Attention(Q,K,V)Softmax(\frac{QK^T}{\sqrt{d}})V $$其中$Q$和$K^T$的乘积代表输入$x$的注意力为避免该值变得过大通常除以$d$的平方根进行缩放并对每行进行softmax归一化与$V$相乘后得到一个$n*d$的矩阵。参数说明参数名输入/输出描述数据类型数据格式query输入公式中的输入Q。FLOAT16、BFLOAT16、INT8NDkey输入公式中的输入K。FLOAT16、BFLOAT16、INT8、INT4NDvalue输入公式中的输入V。FLOAT16、BFLOAT16、INT8、INT4NDattentionOut输出公式中的输出。FLOAT16、BFLOAT16、INT8ND约束说明该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。入参为空的处理算子内部需要判断参数query是否为空如果是空则直接返回。参数query不为空Tensor参数key、value为空tensor(即S2为0)则attentionOut填充为全零。attentionOut为空Tensor时AscendCLNN框架会处理。其余在上述参数说明中标注了“可传入nullptr”的入参为空指针时不进行处理。参数key、value中对应Tensor的shape需要完全一致非连续场景下 key、value的tensorlist中的batch只能为1个数等于query的BN和D需要相等。由于tensorlist限制, 非连续场景下B不能大于256。当Q_S大于1时querykeyvalue输入功能使用限制如下支持B轴小于等于65536。如果输入类型为INT8且D轴不是32字节对齐则B轴的最大支持值为128。若输入类型为FLOAT16或BFLOAT16且D轴不是16字节对齐B轴同样仅支持到128。支持N轴小于等于256支持D轴小于等于512。inputLayout为BSH或者BSND时建议N*D小于65535。S支持小于等于2097152020M。部分长序列场景下如果计算量过大可能会导致pfa算子执行超时aicore error类型报错errorStr为:timeout or trap error此场景下建议做S切分处理注这里计算量会受B、S、N、D等的影响值越大计算量越大。典型的会超时的长序列即B、S、N、D的乘积较大场景包括但不限于BQ_NQ_SDKV_NKV_S1202097152256120971521220971520256220971520201209715225612097152110209715251212097152D轴限制query、key、value或attentionOut类型包含INT8时D轴需要32对齐query、key、value或attentionOut类型包含INT4时D轴需要64对齐类型全为FLOAT16、BFLOAT16时D轴需16对齐。当Q_S等于1时querykeyvalue输入功能使用限制如下支持B轴小于等于65536支持N轴小于等于256支持D轴小于等于512。query、key、value输入类型均为INT8的场景暂不支持。在INT4伪量化场景下aclnn单算子调用支持KV INT4输入或者INT4拼接成INT32输入建议通过dynamicQuant生成INT4格式的数据因为dynamicQuant就是一个INT32包括8个INT4。在INT4伪量化场景下若KV INT4拼接成INT32输入那么KV的N、D或者H是实际值的八分之一prefix同理。key、value在特定数据类型下存在对于D轴的限制key、value输入类型为INT4INT32时D轴需要64对齐INT32仅支持D 8对齐。query、key、value数据排布格式支持从多种维度解读其中BBatch表示输入样本批量大小、SSeq-Length表示输入样本序列长度、HHead-Size表示隐藏层的大小、NHead-Num表示多头数、DHead-Dim表示隐藏层最小的单元尺寸且满足DH/N、T表示所有Batch输入样本序列长度的累加和。调用说明调用方式样例代码说明aclnn接口test_aclnn_FusedInferAttentionScoreV4通过aclnnFusedInferAttentionScoreV4调用PromptFlashAttentionV3算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

相关新闻