开启支持输入转NZ格式功能

key

[num_tokens, num_head, head_size_k]

int8/float16/bf16

ND

head_size_k*sizeof(dtype_k)满足32字节对齐，同时小于98304。

value

[num_tokens, num_head, head_size_v]

float16/bf16

ND

head_size_v*sizeof(dtype_v)满足32字节对齐，同时小于98304。key支持int8时，value可以是float16或者bf16。

keyCache

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

int8/float16/bf16

NZ

当数据类型为int8时，last_dim_k为32，当数据类型是float16/bf16时，last_dim_k为16。数据类型与key保持一致。

valueCache

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

float16/bf16

NZ

last_dim_v只能是16。数据类型与value保持一致。

slotMapping

[num_tokens]

int32

ND

last_dim_v只能是16。值域范围为(INT32_MIN, num_blocks * block_size)且不存在重复数值。取值为负数时不进行任何搬运操作。

keyCacheOut

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

int8/float16/bf16

NZ

当数据类型为int8时，last_dim_k为32，当数据类型是float16/bf16时，last_dim_k为16。

valueCacheOut

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

float16/bf16

NZ

last_dim_v只能是16。