Enabling the Single-input Single-output Function

key

[num_tokens, num_head, k_head_size]

float16/bf16

ND

Key of multiple tokens in the current step.

keyCache

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

All key caches.

slotMapping

[num_tokens]

int32

ND

• Storage offset of each token key or value in the cache, that is, (block_id × block_size + offset_in_block).
• The value range is [0, num_blocks × kv_heads × block_size), and the value must be unique. That is, num_tokens <= num_blocks × block_size.

kStrides

[3]

int64

ND

Stride of the k tensor. The last element of kStrides is 1. Stride for jumping to each dimension of the next element of the k tensor. The default value is NULL. The number of dimensions is the same as that of the input tensor.

kOffset

[1]

int64

ND

Offset of the k tensor relative to the start address of the memory. The default value is 0.

keyCacheOut

[num_blocks, block_size, k_num_head, k_head_size]

float16/bf16

ND

All key caches. The address is the same as that of the input keyCache.