开启单入单出功能

key

[num_tokens, k_num_head, k_head_size]

float16/bf16/int8

ND

当前step多个token的key。

keyCache

[num_blocks, block_size, k_num_head, k_head_size]

float16/bf16/int8

ND

当前layer的key cache。

slotMapping

[num_tokens]

int32

ND

每个token key在cache中的存储偏移，即（block_id * block_size + offset_in_block），值域范围为(INT32_MIN, num_blocks * block_size)且不存在重复数值。取值为负数时不进行任何搬运操作。

keyCacheOut

[num_blocks, block_size, k_num_head, k_head_size]

float16/bf16/int8

ND

所有的key cache。与输入keyCache为同一地址。