Disabling the Multi-Head Compression Function

key

[num_tokens, num_head, k_head_size]

float16/bf16

ND

Key of multiple tokens in the current step. If the multi-head compression function is disabled, the last dimension of the key can be different from that of value.

value

[num_tokens, num_head, v_head_size]

float16/bf16

ND

Value of multiple tokens in the current step. If the multi-head compression function is disabled, the last dimension of the key can be different from that of value.

keyCache

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

Key cache of the current layer. When multi-head compression is disabled, the last dimension of keyCache can be different from that of valueCache.

valueCache

[num_blocks, block_size, num_head, v_head_size]

float16/bf16

ND

Value cache of the current layer. When multi-head compression is disabled, the last dimension of keyCache can be different from that of valueCache.

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 (INT32_MIN, num_blocks × block_size) and no duplicate value exists. If the value is a negative number, no movement operation is performed.

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 key.

vStrides

[3]

int64

ND

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

kOffset

[1]

int64

ND

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

vOffset

[1]

int64

ND

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

keyCacheOut

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

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

valueCacheOut

[num_blocks, block_size, num_head, v_head_size]

float16/bf16

ND

All value caches. The address is the same as that of valueCache.