不开启多头压缩功能

compressType=COMPRESS_TYPE_UNDEFINED

输入

参数

维度

数据类型

格式

描述

key

[num_tokens, num_head, k_head_size]

float16/bf16

ND

当前step多个token的key,不开启多头压缩功能时支持key和value最后一维不同。

value

[num_tokens, num_head, v_head_size]

float16/bf16

ND

当前step多个token的value,不开启多头压缩功能时支持key和value最后一维不同。

keyCache

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

当前layer的key cache,不开启多头压缩功能时支持keyCache和valueCache最后一维不同。

valueCache

[num_blocks, block_size, num_head, v_head_size]

float16/bf16

ND

当前layer的value cache,不开启多头压缩功能时支持keyCache和valueCache最后一维不同。

slotMapping

[num_tokens]

int32

ND

每个token key或value在cache中的存储偏移,即(block_id * block_size + offset_in_block),值域范围为(INT32_MIN, num_blocks * block_size)且不存在重复数值。取值为负数时不进行任何搬运操作。

kStrides

[3]

int64

ND

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

vStrides

[3]

int64

ND

v tensor 步幅。vStrides最后一个元素等于1。默认NULL,跳转到v tensor的下一个元素的各个维度的步幅。维度数与intensor相同。

kOffset

[1]

int64

ND

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

vOffset

[1]

int64

ND

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

输出

参数

维度

数据类型

格式

描述

keyCacheOut

[num_blocks, block_size, num_head, k_head_size]

float16/bf16

ND

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

valueCacheOut

[num_blocks, block_size, num_head, v_head_size]

float16/bf16

ND

所有的value cache。与输入valueCache为同一地址。