开启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

• 每个head每个batch的key或value在cache中的存储偏移，即（block_id * block_size + offset_in_block）。
• 取值范围为[0, num_blocks * kv_heads * block_size)，且不能存在重复值。

  即意味着num_tokens <= num_blocks*block_size。

wins

[batch*num_head]

int32

ND

压缩量。wins[i] < seqLens[floor(i/num_head]，且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为同一地址。