ReshapeAndCacheOperation

功能

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

图1 ReshapeAndCacheOpeartion功能示意图

算子上下文

图2 ReshapeAndCacheOpeartion算子上下文

如上图所示，为典型rope场景下pa decoder layer的部分结构。ReshapeAndCache与paged attention配合使用，用于存储某layer上各个token的历史key和value。与加速库的KVCache算子不同，ReshapeAndCache算子只维护某一层的kvcache而非一个服务的全部kvcache。因此，若模型中有n个layer，使用ReshapeAndCache的情况下需要n个k_cache tensor与n个v_cache tensor管理历史kv。

定义

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

参数列表

输入输出

• 不开启多头压缩功能

  compressType=COMPRESS_TYPE_UNDEFINED

  • 输入

    参数	维度	数据类型	格式	描述
    key	[num_tokens, num_head, k_head_size]	float16/bf16/int8	ND	当前step多个token的key，不开启多头压缩功能时支持key和value最后一维不同。
    value	[num_tokens, num_head, v_head_size]	float16/bf16/int8	ND	当前step多个token的value，不开启多头压缩功能时支持key和value最后一维不同。
    keyCache	[num_blocks, block_size, num_head, k_head_size]	float16/bf16/int8	ND/NZ	当前layer的key cache，不开启多头压缩功能时支持keyCache和valueCache最后一维不同。
    valueCache	[num_blocks, block_size, num_head, v_head_size]	float16/bf16/int8	ND/NZ	当前layer的value cache，不开启多头压缩功能时支持keyCache和valueCache最后一维不同。
    slotMapping	[num_tokens]	int32	ND	每个token key或value在cache中的存储偏移，即（block_id * block_size + offset_in_block），值域范围为(-num_blocks * block_size, num_blocks * block_size)且不存在重复数值。

  • 输出

    参数	维度	数据类型	格式	描述
    keyCacheOut	[num_blocks, block_size, num_head, k_head_size]	float16/bf16/int8	ND/NZ	所有的key cache。与输入keyCache为同一地址。
    valueCacheOut	[num_blocks, block_size, num_head, v_head_size]	float16/bf16/int8	ND/NZ	所有的value cache。与输入valueCache为同一地址。

• 开启alibi场景下的多头压缩功能

  compressType=COMPRESS_TYPE_KVHEAD

  • 输入

    参数	维度	数据类型	格式	描述
    key	[num_tokens, num_head, head_size]	float16/bf16/int8	ND	当前step多个token的key。最后一维需要是32的倍数。
    value	[num_tokens, num_head, head_size]	float16/bf16/int8	ND	当前step多个token的value。最后一维需要是32的倍数。
    keyCache	[num_blocks, block_size, 1, head_size]	float16/bf16/int8	ND	所有的key cache。
    valueCache	[num_blocks, block_size, 1, head_size]	float16/bf16/int8	ND	所有的value cache。
    slotMapping	[batch*num_head]	int32	ND	每个token key或value在cache中的存储偏移，即（block_id * block_size + offset_in_block）。 取值范围为[0, num_blocks * kv_heads * block_size)，且不能存在重复值。
    wins	[batch*num_head]	int32	ND	压缩量。wins[i] < seqLens[i]，且wins内的值需要大于等于0，值为0时不压缩。
    seqLens	[batch]	int32	ND	每个batch的实际seqLen，内部值都大于0。

  • 输出

    参数	维度	数据类型	格式	描述
    keyCacheOut	[num_blocks, block_size, 1, head_size]	float16/bf16/int8	ND	所有的key cache。与输入keyCache为同一地址。
    valueCacheOut	[num_blocks, block_size, 1, head_size]	float16/bf16/int8	ND	所有的value cache。与输入valueCache为同一地址。

• 开启rope场景下的多头压缩功能

  compressType=COMPRESS_TYPE_KVHEAD_ROPE

  • 输入

    参数	维度	数据类型	格式	描述
    key	[num_tokens, num_head, head_size]	float16/bf16	ND	当前step多个token的key。最后一维需要是32的倍数。
    value	[num_tokens, num_head, head_size]	float16/bf16	ND	当前step多个token的value。最后一维需要是32的倍数。
    keyCache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	所有的key cache。
    valueCache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	所有的value cache。
    slotMapping	[batch*num_head]	int32	ND	每个token key或value在cache中的存储偏移，即（block_id * block_size + offset_in_block）。 取值范围为[0, num_blocks * kv_heads * block_size)，且不能存在重复值。
    wins	[batch*num_head]	int32	ND	压缩量。wins[i] < seqLens[i]，且wins内的值需要大于等于0，值为0时不压缩。
    seqLens	[batch]	int32	ND	每个batch的实际seqLen，内部值都大于0。
    offsetIndex	[batch*num_head]	int32	ND	每个batch每个head的压缩起点。取值范围为[-1, seqLens[i]-wins[i]]；取值为-1时不进行压缩操作，wins的值无要求。

  • 输出

    参数	维度	数据类型	格式	描述
    keyCacheOut	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	所有的key cache。与输入keyCache为同一地址。
    valueCacheOut	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	所有的value cache。与输入valueCache为同一地址。

规格约束

• ntokens = sum(seqlen[i])，i=0、1...batch-1。
• Atlas 推理系列产品 上不支持压缩功能，且数据格式标明有NZ的需为NZ格式，keyCache和valueCache的维度需为[num_blocks, num_head*head_size/16, block_size, 16]，其中最后一维必须为16。
• 在key和value最后一维不同的情况下只支持 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品 ，此场景下keyCache、valueCache、keyCacheOut、valueCacheOut只支持ND格式。
• key、value、keyCache、valueCache、keyCacheOut 、valueCacheOut的数据类型都必须完全一致。
• keyCache、valueCache、和keyCacheOut 、valueCacheOut的数据格式必须一致。
• 非压缩场景下num_tokens <= num_blocks*block_size，压缩场景下cache的大小要大于压缩后的大小。