Q/K/V 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).
Online full quantization requires four input tensors: qkDescale, qkOffset, vpvDescale, and vpvOffset. When offline quantization is used, pScale needs to be passed as the input tensor based on the new input tensor for online full quantization.
The following table lists the dimensions that need to be specified for input tensors.
Parameter |
Dimension |
Data Type |
Format |
npu or cpu |
Description |
|---|---|---|---|---|---|
query |
[nTokens, headNum, headSize] |
int8 |
ND |
npu |
Query matrix. |
key |
[nTokens, headNum, headSize] |
int8 |
ND |
npu |
Key matrix. |
value |
[nTokens, headNum, headSize] |
int8 |
ND |
npu |
Value matrix. |
mask |
Same as Mask Types. Different from FA when mask compression is enabled. See mask. |
float16/bf16 |
ND/NZ |
npu |
The four shapes correspond to: 1. Square matrix with all batches being the same. 2. Matrix with different batches. 3. Vectors with different batches. 4. ALiBi use cases. This tensor is not passed when maskType is undefined. |
seqLen |
[batch] |
int32 |
ND |
cpu |
1: incremental or full mode. > 1: full mode. |
slopes |
[headNum] |
ND |
npu |
When maskType is set to ALiBi-compressed and the shape of mask is set to [256,256], this tensor needs to be passed. It is the coefficient of each head of the ALiBi mask. |
|
qkDescale |
[headNum] |
float |
ND |
npu |
Dequantization scale of Q x K^T. |
qkOffset |
[headNum] |
int32 |
ND |
npu |
Dequantization offset of Q*K^T. Reserved. A non-null tensor needs to be passed. This parameter is not used currently. |
vpvDescale |
[headNum] |
float |
ND |
npu |
When quantType is 2, this parameter indicates the dequantized scale of P x V. When quantType is 3, this parameter indicates the dequantized scale of V. |
vpvOffset |
[headNum] |
int32 |
ND |
npu |
This tensor is passed only for full quantization (quantType=2 or 3).
Reserved. A non-null tensor needs to be passed. This parameter is not used currently. |
pScale |
[headNum] |
float |
ND |
npu |
Offline quantization scale of P.
|
output |
[nTokens, headNum, headSize] |
float16/bf16 |
ND |
npu |
Output. |
pScale is passed only for offline quantization.
Constraints
When using full quantization, specify the data type of the output tensor by using the outDataType parameter, which can only be ACL_FLOAT16 or ACL_BF16.
calcType must be set to PA_ENCODER.
Not supported by