ScatterPaKvCacheWithKScale【免费下载链接】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 训练系列产品×功能说明算子功能更新KvCache中指定位置的key和value同时更新key的scale值。输入输出支持以下场景场景一key:[batch * seq_len, num_head, head_size] value:[batch * seq_len, num_head, head_size] key_cache:[num_blocks, num_head, block_size, head_size] value_cache:[num_blocks, num_head, block_size, head_size] slot_mapping:[batch * seq_len] key_scale:[batch * seq_len, num_head] key_scale_cache:[num_blocks, num_head, block_size, 1] cache_layout:BNBD其中key和value的dtype为FLOAT8_E5M2或FLOAT8_E4M3FNkey_scale和key_scale_cache的dtype为FLOAT。计算公式对于每个tokeni ∈ [0, num_tokens)和每个头j ∈ [0, num_head)block_idx slot_mapping[i] // block_size block_offset slot_mapping[i] % block_size key_cache[block_idx][j][block_offset][:] key[i][j][:] value_cache[block_idx][j][block_offset][:] value[i][j][:] key_scale_cache[block_idx][j][block_offset][0] key_scale[i][j]其中num_tokens batch * seq_lenblock_idxslot_mapping映射到的block索引block_offsetblock内的偏移量Ascend 950PR/Ascend 950DT 仅支持场景一。参数说明参数名输入/输出/属性描述数据类型数据格式key输入待更新的key值当前step多个token的key。FLOAT8_E5M2、FLOAT8_E4M3FNNDvalue输入待更新的value值当前step多个token的value。FLOAT8_E5M2、FLOAT8_E4M3FNNDkey_cache输入/输出需要更新的key cache当前layer的key cache。与key保持一致NDvalue_cache输入/输出需要更新的value cache当前layer的value cache。与value保持一致NDslot_mapping输入每个token key或value在cache中的存储偏移。INT32、INT64NDkey_scale输入待更新的key scale值当前step多个token的key scale。FLOATNDkey_scale_cache输入/输出需要更新的key scale cache当前layer的key scale cache。FLOATNDcache_layout属性表示key_cache和value_cache的内存排布格式。当传空指针或BNBD时表示格式为[num_blocks, num_head, block_size, head_size]。STRING-约束说明确定性计算默认确定性实现。key、value、key_cache、value_cache的数据类型必须一致slot_mapping的取值范围[0, num_blocks*block_size-1]且slot_mapping内的元素值保证不重复重复时不保证正确性key和value的前两维shape必须相同key_scale是两维tensorshape为[batch * seq_len, num_head]尾轴可以不连续key_scale_cache是四维tensorshape为[num_blocks, num_head, block_size, 1]最后一维必须为1尾轴必须连续。调用说明调用方式样例代码说明aclnn接口test_aclnn_ScatterPaKvCacheWithKScale通过aclnnScatterPaKvCacheWithKScale调用ScatterPaKvCacheWithKScale算子图模式test_geir_ScatterPaKvCacheWithKScale通过算子IR调用ScatterPaKvCacheWithKScale算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考