QKV Full Quantization
Description
Quantized q, k, and v can be passed to reduce the graphics memory usage.
How to Enable
Set quantType to TYPE_QUANT_QKV_OFFLINE (offline quantization) or TYPE_QUANT_QKV_ONLINE (online quantization).
Reuse kDescale and vDescale . When offline quantization is used, pass pScale as the input tensor.
The required dimensions for input tensors are as follows.
Parameter |
Dimension |
Data Type |
Format |
cpu or npu |
|---|---|---|---|---|
query |
[num_tokens, num_head, head_size] |
int8 |
ND |
npu |
keyCache |
[num_blocks, block_size, kv_head_num, head_size] |
int8 |
ND |
npu |
valueCache |
[num_blocks, block_size, kv_head_num, head_size] |
int8 |
ND |
npu |
blockTables |
[num_tokens, max_num_blocks_per_query] |
int32 |
ND |
npu |
contextLens |
[batch] |
int32 |
ND |
cpu |
kDescale |
[head_num] |
float |
ND |
npu |
vDescale |
[head_num] |
float |
ND |
npu |
pScale |
[head_num] |
float |
ND |
npu |
attnOut |
[n_tokens, num_head, head_size] |
float16/bf16 |
ND |
npu |
pScale is passed only in the offline quantization scenario.
Constraints
When full quantization is used, specify the data type of the output tensor by using the outDataType parameter, which can only be ACL_FLOAT16 or ACL_BF16.
Not supported by the