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).

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 when mask compression is enabled. For details, see compressing mask.	float16/bf16	ND/NZ	npu	The four shapes are described as follows: 1. All batches are the same, square matrix. 2. Matrix with different batches. 3. Vectors with different batches. 4. Alibi. This tensor is not passed when maskType is undefined.
seqLen	[batch]	int32	ND	cpu	1: incremental or full mode. > 1: full mode.
slopes	[headNum]	Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products: float16	ND	npu	When maskType is set to ALiBi compression 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 in full quantization scenarios (quantType=2 or 3). When quantType is 2, this parameter indicates the dequantization offset of P*V. When quantType is 3, this parameter indicates the dequantization offset of V. 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. This tensor needs to be passed when offline full quantization is enabled (quantType = 2). This tensor is not passed when online full quantization is enabled (quantType=3).
output	[nTokens, headNum, headSize]	float16/bf16	ND	npu	Output.

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.

calctype must be set to PA_ENCODER.

Not supported by Atlas inference products.

Parent topic: Functions