开启单入单出功能

key

[num_tokens, num_head, k_head_size]

float16/bf16

ND

当前step多个token的key。

keyCache

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

所有的key cache。

slotMapping

[num_tokens]

int32

ND

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

kStrides

[3]

int64

ND

k tensor 步幅。kStrides最后一个元素等于1。默认NULL，跳转到k tensor的下一个元素的各个维度的步幅。维度数与输入tensor相同。

kOffset

[1]

int64

ND

默认0，k tensor相对于内存首地址的偏移。

keyCacheOut

[num_blocks, block_size, k_num_head, k_head_size]

float16/bf16

ND

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