ReshapeAndCacheOperation

功能

遍历每个key和value，将key和value(num_heads, head_size)以slotmapping填入key_cache/value_cache指定位置。

约束

输出keyCache，valueCache与同名输入指向同一地址，即进行原地修改。
在key和value最后一维不同的情况下只支持Atlas 800I A2推理产品，即keyCache, valueCache, keyCacheOut, valueCacheOut只支持ND格式。
key、value、keyCache、valueCache、keyCacheOut 、valueCacheOut的数据类型都必须完全一致。
keyCache、valueCache、和keyCacheOut 、valueCacheOut的数据格式必须一致。

定义

struct ReshapeAndCacheParam {
   enum CompressType : int {
      COMPRESS_TYPE_UNDEFINED = 0,
      COMPRESS_TYPE_KVHEAD
   };
   CompressType compressType = COMPRESS_TYPE_UNDEFINED;
};

成员

成员名称	描述
compressType	压缩类型。 COMPRESS_TYPE_UNDEFINED：默认值，不压缩。 COMPRESS_TYPE_KVHEAD：压缩key_cache, value_cache的kvHead维度。

成员名称

描述

compressType

压缩类型。

COMPRESS_TYPE_UNDEFINED：默认值，不压缩。
COMPRESS_TYPE_KVHEAD：压缩key_cache, value_cache的kvHead维度。

输入

参数	维度	数据类型	格式	描述
key	[num_tokens, k_num_head, head_size]	float16/bfloat16/int8	ND	当前step多个token的key。
value	[num_tokens, v_ num_head, head_size]	float16/bfloat16/int8	ND	当前step多个token的value。
keyCache	[num_blocks, block_size, k_num_head, head_size]	float16/bfloat16/int8	ND/NZ	当前layer的key cache。
valueCache	[num_blocks, block_size, v_ num_head, head_size]	float16/bfloat16/int8	ND/NZ	当前layer的value cache。
slotMapping	[num_tokens]	int32	ND	每个token key或value在cache中的存储偏移，即（block_id * block_size + offset_in_block）。
wins	[num_tokens]	int32	ND	当开启多头压缩功能时需传入此tensor。 Atlas 推理系列产品不支持此功能。
seqLen	[batch]	int32	ND	当开启多头压缩功能时需传入此tensor。 Atlas 推理系列产品不支持此功能。

输出

参数	维度	数据类型	格式	描述
keyCacheOut	[num_blocks, block_size, k_head_num, head_size]	float16/bfloat16/int8	ND/NZ	所有的key cache。
valueCacheOut	[num_blocks, block_size, v_num_head, head_size]	float16/bfloat16/int8	ND/NZ	所有的value cache。

父主题： atb/infer_op_params.h