多头自适应压缩(alibi场景)

功能

alibi mask场景下,对kv的 head_num维度进行压缩,提高内存利用率,精度不变。

又名Razor Attention。

开启方式

参数“compressType”置为COMPRESS_TYPE_KVHEAD。

若干输入tensor需要传入指定的维度,如下:

参数

维度

数据类型

格式

cpu or npu

query

[num_tokens * kv_head_num, num_head / kv_head_num, head_size]

float16/bf16

ND

npu

keyCache

[num_blocks, block_size, 1, head_size]

float16/bf16/int8

ND

npu

valueCache

[num_blocks, block_size, 1, head_size]

float16/bf16/int8

ND

npu

blockTables

[num_tokens * kv_head_num, max_num_blocks_per_query]

int32

ND

npu

contextLens

[batch * num_head]

int32

ND

cpu

attnOut

[n_tokens, num_head, head_size]

float16/bf16

ND

npu

特殊约束

不支持Atlas 推理系列产品

需与ReshapeAndCache算子的相应功能配合使用。