CANN稀疏注意力元数据算子文档

SparseAttnSharedkvMetadata【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 推理系列产品√功能说明API功能SparseAttnSharedkvMetadata算子旨在生成一个任务列表包含每个AIcore的Attention计算任务的起止点的Batch、Head、以及 Q 和 K 的分块的索引供后续SparseAttnSharedkv算子使用。函数原型custom.npu_sparse_attn_sharedkv_metadata(num_heads_q, num_heads_kv, head_dim, *, cu_seqlens_qNone, cu_seqlens_ori_kvNone, cu_seqlens_cmp_kvNone, seqused_qNone, seqused_kvNone, batch_size0, max_seqlen_q0, max_seqlen_kv0, ori_topk0, cmp_topk0, cmp_ratio-1, ori_mask_mode4, cmp_mask_mode3, ori_win_left127, ori_win_right0, layout_qBSND, layout_kvPA_ND, has_ori_kvTrue, has_cmp_kvTrue, devicenpu:0) - TensorTransformer SparseAttnSharedkvMetadata 算子实现参考: SparseAttnSharedkvMetadata参数说明num_heads_qint必选参数表示公式中的$Q$的多头数目前仅支持64。num_heads_kvint必选参数表示公式中的$\tilde{K}$和$\tilde{V}$的多头数目前仅支持1。head_dimint必选参数表示注意力头的维度。*必选参数代表其之前的变量是位置相关的必须按照顺序输入之后的变量是可选参数位置无关需要使用键值对赋值不赋值会使用默认值。cu_seqlens_qTensor可选参数当layout_query为TND时表示不同Batch中q的有效token数维度为B1大小为参数中每个元素的值表示目前batch与之前所有batch的token数总和即前缀和数据类型支持int32。cu_seqlens_ori_kvTensor可选参数当layout_kv为TND时表示不同Batch中ori_kv的有效token数维度为B1大小为参数中每个元素的值表示目前batch与之前所有batch的token数总和即前缀和数据类型支持int32。目前layout_kv仅支持PA_ND故设置此参数无效。cu_seqlens_cmp_kvTensor可选参数当layout_kv为TND时表示不同Batch中cmp_kv的有效token数维度为B1大小为参数中每个元素的值表示目前batch与之前所有batch的token数总和即前缀和数据类型支持int32。目前layout_kv仅支持PA_ND故设置此参数无效。seqused_qTensor可选参数表示不同Batch中q实际参与运算的token数维度为B数据格式支持ND数据类型支持int32不输入则所有token均参与运算。目前暂不支持指定该参数。seqused_kvTensor可选参数表示不同Batch中ori_kv实际参与运算的token数维度为B数据格式支持ND数据类型支持int32不输入则所有token均参与运算。batch_sizeint可选参数表示输入样本批量大小默认值为None。max_seqlen_qint可选参数表示所有batch中q的最大有效token数。max_seqlen_kvint可选参数表示所有batch中ori_kv的最大有效token数。ori_topkint可选参数表示通过QLI算法从ori_kv中筛选出的关键稀疏token的个数。目前暂不支持指定该参数默认值为None。cmp_topkint可选参数表示通过QLI算法从cmp_kv中筛选出的关键稀疏token的个数目前仅支持512默认值为None。cmp_ratioint可选参数表示对ori_kv的压缩率数据范围支持4/128默认值为None。ori_mask_modeint可选参数表示q和ori_kv计算的mask模式目前仅支持输入默认值4代表band模式的mask。cmp_mask_modeint可选参数表示q和cmp_kv计算的mask模式目前仅支持输入默认值3代表rightDownCausal模式的mask对应以右顶点为划分的下三角场景。ori_win_leftint可选参数表示q和ori_kv计算中q对过去token计算的数量目前仅支持默认值127。ori_win_rightint可选参数表示q和ori_kv计算中q对未来token计算的数量目前仅支持默认值0。layout_qstr可选参数表示输入q的数据排布格式默认值为BSND目前支持传入BSND和TND。layout_kvstr可选参数表示输入ori_kv和cmp_kv的数据排布格式目前仅支持传入默认值PA_NDPageAttention。has_ori_kvbool可选参数表示是否传入ori_kv默认值为true。has_cmp_kvbool可选参数表示是否传入cmp_kv默认值为true。devicestr可选参数用于获取设备信息默认值为None。返回值说明metadataTensor每个cube核上FlashAttention计算任务的Batch、Head、以及 Q 和 K 的分块的索引以及每个vector核上FlashDecode的规约任务索引。约束说明该接口支持推理场景下使用。该接口支持aclgraph模式。Tensor不能全传None。Atlas A3 推理系列产品 调用示例支持单算子模式调用和aclgraph模式调用作为SparseAttnSharedkv算子的前序算子调用示例见SparseAttnSharedkv调用示例。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考