CANN/ops-transformer KV缓存收集算子

GatherPaKvCache【免费下载链接】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 训练系列产品×功能说明算子功能根据blockTables中的blockId值、seqLens中key/value的seqLen从keyCache/valueCache中将内存不连续的token搬运、拼接成连续的key/value序列。计算逻辑keyRef/valueRef的第一个维度取决于seq_lens大小。如果isSeqLensCumsum为true则seqLens中最后一个值即为keyRef/valueRef的第一个维度大小 keyRef[dim0] seqLens[-1]如果isSeqLensCumsum为false则seqLens中所有值的累加和即为keyRef/valueRef的第一个维度大小keyRef[dim0] sum(seqLens)关于keyRef、valueRef的一些限制条件如下每个token大小控制在148k以内例如对于fp16/bf16类型num_heads * head_size(keyRef/valueRef)取128*576。示例keyCache_shape: [128, 128, 16, 144] valueCache_shape: [128, 128, 16, 128] blockTables_shape: [16, 12] seqLens_shape: [16] keyRef_shape: [8931, 16, 144] valueRef_shape: [8931, 16, 128] seqOffset_shape: [16] out1_shape: [8931, 16, 144] out2_shape: [8931, 16, 128]参数说明参数名输入/输出/属性描述数据类型数据格式keyCache输入当前层存储的key向量缓存INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FNNDvalueCache输入当前层存储的value向量缓存INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FNFRACTAL_NZblockTables输入每个batch中KV Cache的逻辑块到物理块的映射关系INT32、INT64NDseqLens输入每个batch对应的序列长度INT32、INT64NDkeyRef输入/输出当前层的key向量INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FNNDvalueRef输入/输出当前层的value向量INT8, FLOAT16, BFLOAT16, FLOAT, UINT8, INT16, UINT16, INT32, UINT32, HIFLOAT8, FLOAT8_E5M2, FLOAT8_E4M3FNNDseqOffset输入blockTables获取blockId时存在的首偏移INT32、INT64NDcacheMode输入表示输入的数据排布格式支持Norm、PA_NZStringNDisSeqLensCumsum输入表示seqLens是否为累加和。false表示非累加和BOOLND约束说明无调用说明调用方式样例代码说明aclnn接口test_aclnn_gather_pa_kv_cache通过aclnnGatherPaKvCache调用GatherPaKvCache算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考