Multi-head Adaptive Compression (alibi)
Description
In the alibi mask scenario, compress the head_num dimension of the KV to improve the memory usage without changing the precision.
Also called Razor Attention.
How to Enable
Set compressType to COMPRESS_TYPE_KVHEAD.
The required dimensions for input tensors are as follows.
Parameter |
Dimension |
Data Type |
Format |
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 |
Constraints
Not supported by the
Must be used together with the ReshapeAndCache operator.