接口功能:
[object Object]算子实现基于共享KV(Key=Value)的稀疏注意力计算,支持SWA(Sliding Window Attention)、CSA(Compressed Sparse Attention)、HCA(Heavily Compressed Attention)三类Attention计算场景。与[object Object]的区别在于,本算子支持KV的per-token-group量化输入。该算子适用于大语言模型推理场景,通过滑动窗口和KV压缩机制大幅降低长序列注意力计算的开销。调用时需要使用[object Object]生成的任务列表[object Object]。该算子不建议单独使用,建议与aclnnMixedQuantSparseFlashMlaMetadata算子配合使用,形成完整的工作流。
典型调用流程如下:
- 准备
[object Object]、[object Object]、[object Object]、序列长度、[object Object]、[object Object]等输入。 - 调用
[object Object]生成[object Object]。 - 调用
[object Object],将上一步得到的[object Object]传入主算子。
- 准备
计算公式:
其中(共享KV),由滑动窗口内的原始KV和因果边界内的压缩KV拼接而成,具体参与计算的KV范围由模板模式和mask参数决定:
滑动窗口部分(oriKv):对第个Query token,其因果对角线位置为,窗口范围为。
压缩KV部分(cmpKv):因果边界阈值为。HCA场景取内的连续压缩KV;CSA场景通过TopK索引从压缩KV中按需收集,仅保留的块。
注意力计算采用Online Softmax(Flash Attention V2),S2方向按512分块循环,sinks作为每行softmax的初始最大值:
符号说明
[object Object]undefined
每个算子分为,必须先调用[object Object]接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用[object Object]执行实际计算。
参数说明
[object Object]- [object Object]Ascend 950PR/Ascend 950DT[object Object]:N1/N2支持2、4、8、16、32、64、128,不支持1。
返回值
第一段接口完成入参校验,出现以下场景时报错:
[object Object]
确定性计算
- aclnnMixedQuantSparseFlashMla默认采用确定性实现,相同输入多次调用结果一致。
使用约束
- 本算子仅支持推理场景。
- 除
[object Object]和[object Object]等预留输入可传入nullptr或空Tensor外,其余已传入Tensor不支持为空。 [object Object]参数必须传入,由[object Object]算子生成,shape固定为(1024,)。[object Object]为主算子和[object Object]的可选入参;传入后用于按[object Object]恢复cmp侧mask使用的压缩前长度。[object Object]仅支持64。[object Object]仅支持4,[object Object]仅支持3,[object Object]仅支持127,[object Object]仅支持0。
三种Attention场景输入要求
[object Object]undefined
[object Object]约束:SWA场景仅支持1。Layout约束
[object Object]和[object Object]组合仅支持"BSND"/"BSND"、"TND"/"TND"、"BSND"/"PA_BBND"、"TND"/"PA_BBND";非PA_BBND场景下[object Object]和[object Object]必须一致。- 当
[object Object]为TND时,[object Object]必须传入。 - 当
[object Object]为PA_BBND时,[object Object]必须传入,[object Object]必须传入。BSND场景可选传入[object Object]覆盖每个batch的oriKv有效长度;TND场景使用[object Object]表达oriKv序列边界。 - 当
[object Object]为TND时,[object Object]必须传入。 - 当
[object Object]为TND且存在[object Object]时,[object Object]必须传入。 [object Object]为所有layoutKvOptional下的可选输入,显式传入时用于覆盖cmp侧逻辑有效长度。