Enabling the Function of Converting the Format of the Input to NZ

kvCacheCfg=K_CACHE_V_CACHE_NZ, compressType=COMPRESS_TYPE_UNDEFINED

Input

Parameter

Dimension

Data Type

Format

Description

key

[num_tokens, num_head, k_head_size]

int8/float16/bf16

ND

k_head_size × sizeof(dtype_k) is 32-byte aligned and less than 98304.

value

[num_tokens, num_head, v_head_size]

float16/bf16

ND

v_head_size × sizeof(dtype_v) is 32-byte aligned and less than 98304.

When the data type of key is int8, the data type of value can be float16 or bf16.

keyCache

[num_blocks, num_head*k_head_size // last_dim_k, block_size, last_dim_k]

int8/float16/bf16

NZ

When the data type is int8, last_dim_k is 32.

When the data type is float16/bf16, last_dim_k is 16.

The data type is the same as that of key.

block_size must be 16-aligned.

valueCache

[num_blocks, num_head*v_head_size // last_dim_v, block_size, last_dim_v]

float16/bf16

NZ

last_dim_v can only be 16.

The data type is the same as that of value.

block_size must be 16-aligned.

slotMapping

[num_tokens]

int32

ND

last_dim_v can only be 16.

The value range is (INT32_MIN, num_blocks × block_size), and the value must be unique.

If the value is a negative number, no movement operation is performed.

Output

Parameter

Dimension

Data Type

Format

Description

keyCacheOut

[num_blocks, num_head*k_head_size // last_dim_k, block_size, last_dim_k]

int8/float16/bf16

NZ

When the data type is int8, last_dim_k is 32. When the data type is float16/bf16, last_dim_k is 16.

valueCacheOut

[num_blocks, num_head*k_head_size // last_dim_v, block_size, last_dim_v]

float16/bf16

NZ

last_dim_v can only be 16.