SparseCompressedFlashAttention【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A2 推理系列产品√Atlas A3 推理系列产品√Ascend 950PR/Ascend 950DT 推理系列产品√功能说明API功能SparseCompressedFlashAttention算子旨在完成以下公式描述的Attention计算支持Sparse Compressed Attention。计算公式$$ O \text{softmax}(Q\tilde{K}^T \cdot \text{softmax_scale})\tilde{V} $$其中$\tilde{K}\tilde{V}$为基于ori_kv、cmp_kv以及cmp_kv等入参控制的实际参与计算的 $KV$。函数原型torch.ops.pypto.sparse_compress_flash_attention( query, q_act_seqs, ori_kv, cmp_kv, ori_block_table, cmp_block_table, atten_sink, seqused_kv, cmp_sparse_indices, softmax_scale, win_size, cmp_ratio ) - (Tensor)参数说明queryTensor必选参数对应公式中的$Q$不支持非连续数据格式支持ND数据类型支持bfloat16。shape为[T1*N1,D]其中N1仅支持64。q_act_seqsTensor必选参数在layout_query为TND时生效。表示不同Batch中q的有效token数维度为B1大小为参数中每个元素的值表示当前batch与之前所有batch的token数总和即前缀和因此后一个元素的值必须前一个元素的值数据类型支持int32。ori_kvTensor必选参数对应公式中的$\tilde{K}和\tilde{V}$的一部分为原始不经压缩的KV不支持非连续数据格式支持ND数据类型支持bfloat16layout_kv为PA_ND时shape为[block_num1* ori_block_size, KV_N*D]其中block_num1为PageAttention时block总数ori_block_size为一个block的token数ori_block_size取值为128KV_N仅支持1。cmp_kvTensor必选参数对应公式中的$\tilde{K}和\tilde{V}$的一部分为经过压缩的KV不支持非连续数据格式支持ND数据类型支持bfloat16layout_kv为PA_ND时shape为[block_num2* cmp_block_size, KV_N*D]其中block_num2为PageAttention时block总数cmp_block_size为一个block的token数cmp_block_size取值为128。ori_block_tableTensor必选参数表示PageAttention中oriKvCache存储使用的block映射表。数据格式支持ND数据类型支持int32shape为2维其中第一维长度为B第二维长度不小于所有batch中最大的S2对应的block数量即S2_max / block_size向上取整。cmp_block_tableTensor必选参数表示PageAttention中cmpKvCache存储使用的block映射表。数据格式支持ND数据类型支持int32shape为2维其中第一维长度为B第二维长度不小于所有batch中最大的S3对应的block数量即S3_max / block_size向上取整。atten_sinkTensor必选参数注意力下沉tensor数据格式支持ND数据类型支持float32shape为[N1]。seqused_kvTensor必选参数表示不同Batch中ori_kv实际参与运算的token数维度为B数据格式支持ND数据类型支持int32不输入则所有token均参与运算。cmp_sparse_indicesTensor必选参数代表离散取cmpKvCache的索引不支持非连续数据格式支持ND数据类型支持int32。当layout_query为TND时shape需要传入[Q_T * KV_N, K2]其中K2为对cmp_kv一次离散选取的token数K2仅支持512。softmax_scaledouble必选参数代表缩放系数作为q与ori_kv和cmp_kv矩阵乘后Muls的scalar值数据类型支持float。win_sizeint必选参数窗口大小数据类型支持int32仅支持128。cmp_ratioint必选参数表示对ori_kv的压缩率数据类型支持int数据支持4。返回值说明attention_outTensor公式中的输出。数据格式支持ND数据类型支持bfloat16。shape为[T1*N1,D]。约束说明该接口支持推理场景下使用。该接口支持aclgraph模式。参数q中的D和ori_kv、cmp_kv的D值相等为512。参数q、ori_kv、cmp_kv的数据类型必须保持一致。为了提高算子性能当前q、ori_kv、cmp_kv、attention_out进行了高维合轴处理。仅支持TND格式。block_size支持128。调用方法python3 ops/pypto_python/example/test_sparse_compressed_flash_attention_pypto.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考