MlaPreprocessOperation

Applicable Products

Hardware Model	Supported
Atlas 350 accelerator cards	x
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

Integrates the entire process of PagedAttention input data processing in MLA, including rmsnorm, dequantization, matmul, rope, and reshapeAndCache calculations starting from the hidden state input.

Figure 1 Computation process

Definition

struct MlaPreprocessParam{
    uint32_t wdqDim = 0;
    uint32_t qRopeDim = 0;
    uint32_t kRopeDim = 0;
    float epsilon = 1e-5;
    int32_t qRotaryCoeff = 2;
    int32_t kRotaryCoeff = 2;
    bool transposeWdq = true;
    bool transposeWuq = true;
    bool transposeWuk = true;
    enum CacheMode : uint8_t {
        KVCACHE = 0,
        KROPE_CTKV,
        INT8_NZCACHE,
        NZCACHE
    };
    CacheMode cacheMode = KVCACHE;
    enum QuantMode : uint16_t {
        PER_TENSOR_QUANT_ASYMM = 0,
        PER_TOKEN_QUANT_SYMM,
        PER_TOKEN_QUANT_ASYMM,
        UNQUANT
    };
    QuantMode quantMode = PER_TENSOR_QUANT_ASYMM;
    enum BackendType : int {
        BACKEND_TYPE_ATB = 0,  
        BACKEND_TYPE_ACLNN,    
        BACKEND_TYPE_MAX,      
    };
    BackendType backendType = BACKEND_TYPE_ATB;
    uint8_t rsv[34] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory	Description
wdqDim	uint32_t	0	[0]	No	Split dimension size after Matmul.
qRopeDim	uint32_t	0	[0]	No	Dimension size of q passed to RoPE.
kRopeDim	uint32_t	0	[0]	No	Dimension size of k passed to RoPE.
epsilon	float	1e-5	[1e-5]	No	Added to the denominator to prevent division by 0.
qRotaryCoeff	int32_t	2	2	No	q rotation coefficient.
kRotaryCoeff	int32_t	2	2	No	k rotation coefficient.
transposeWdq	bool	true	(true)	No	Whether to transpose wdq.
transposeWuq	bool	true	(true)	No	Whether to transpose wuq.
transposeWuk	bool	true	(true)	No	Whether to transpose wuk.
cacheMode	uint8_t	0	[0,3]	No	Cache type.
quantMode	uint16_t	0	[0,3]	No	RmsNorm quantization type.
rsv[34]	uint8_t	{0}	[0]	No	Reserved.
Currently, only cacheMode and quantMode take effect.

The user-defined types in the preceding table are described as follows.

cacheMode: type of the input query and kcache. The values are as follows:

KVCACHE: kcache and q are concatenated and then output.
KROPE_CTKV: The input and output kvCache is split into krope and ctkv, and q is split into qrope and qnope.
INT8_NZCACHE: krope and ctkv are converted into the NZ format for output, and ctkv and qnope are quantized into the int8 type through per_head static symmetric quantization.
NZCACHE: krope and ctkv are converted into the NZ format for output.

quantMode: RmsNorm quantization type. The values are as follows:

PER_TENSOR_QUANT_ASYMM: per_tensor static asymmetric quantization, which is the default quantization type.
PER_TOKEN_QUANT_SYMM: per_token dynamic symmetric quantization, which is not implemented.
PER_TOKEN_QUANT_ASYMM: per_token dynamic asymmetric quantization, which is not implemented.
UNQUANT: no quantization, floating-point output, which is not implemented.

Input

Module Function	Parameter	Dimension	Data Type	Format	Description
rmsNormQuant_0	input	[tokenNum, hiddenSize]	float16/bf16	ND	Mandatory.
	gamma0	[hiddenSize]	float16/bf16	ND	Mandatory. Same data type as input.
	beta0	[hiddenSize]	float16/bf16	ND	Mandatory. Same data type as input.
	quantScale0	[1]	float16/bf16	ND	Mandatory. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0. The data type is the same as that of input.
	quantOffset0	[1]	int8	ND	Mandatory. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0.
matmul_0	wdqkv	[2112, hiddenSize]	int8/bf16	NZ	Mandatory. If the data type is the same as that of input and the data type is bf16, rmsNormQuant is disabled.
	deScale0	[2112]	int64/float	ND	Mandatory. If the data type of input is float16, the data type of this parameter is int64. If the data type of input is bf16, the data type of this parameter is float.
	bias0	[2112]	int32	ND	Mandatory. An empty tensor can be passed. This parameter is not passed when quantMode is set to 1 or 3.
rmsNormQuant_1	gamma1	[1536]	float16/bf16	ND	Mandatory. Same data type as input.
	beta1	[1536]	float16/bf16	ND	Mandatory. Same data type as input.
	quantScale1	[1]	float16/bf16	ND	Mandatory. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0. The data type is the same as that of input.
	quantOffset1	[1]	int8	ND	Mandatory. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0.
matmul_1	wuq	[headNum * 192, 1536]	int8/bf16	NZ	Mandatory. If the data type is the same as that of input and the data type is bf16, rmsNormQuant is disabled.
	deScale1	[headNum * 192]	int64/float	ND	Mandatory. If the data type of input is float16, the data type of this parameter is int64. If the data type of input is bf16, the data type of this parameter is float.
	bias1	[headNum * 192]	int32	ND	Mandatory. An empty tensor can be passed. This parameter is not passed when quantMode is set to 1 or 3.
rmsNorm	gamma2	[512]	float16/bf16	ND	Mandatory. Same data type as input.
rope	cos	[tokenNum,64]	float16/bf16	ND	Mandatory. Same data type as input.
rope	sin	[tokenNum,64]	float16/bf16	ND	Mandatory. Same data type as input.
matmulEin	wuk	ND: [headNum, 128, 512] NZ: [headNum, 32, 128, 16]	float16/bf16	ND/NZ	Mandatory. Same data type as input.
reshapeAndCache	kvCache	If cacheMode is 0: [blockNum,blockSize, 1, 576] If cacheMode is 1: [blockNum,blockSize, 1, 512] If cacheMode is 2: [blockNum, 1512/32, blockSize, 32] If cacheMode* is 3: [blockNum, 1*512/16, blockSize, 16]	float16/bf16/int8	ND/NZ	Mandatory. Same data type as input. If cacheMode is set to 1, the tensor shape is split. If cacheMode is set to 2, the format is NZ and the type is int8. If cacheMode is set to 3, the format is NZ.
	kvCacheRope	If cacheMode is 1: [blockNum,blockSize, 1, 64] If cacheMode is 2 or 3: [blockNum, 1*64 / 16 , blockSize, 16]	float16/bf16	ND/NZ	Mandatory. An empty tensor can be passed. This parameter is passed if cacheMode is not set to 0. The data type is the same as that of input. If cacheMode is 2 or 3, the format is NZ.
	slotmapping	[tokenNum]	int32	ND	Mandatory.
quant	ctkvScale	[1]	float16/bf16	ND	Mandatory. An empty tensor can be passed. This parameter is passed if cacheMode is set to 2. The data type is the same as that of input.
quant	qNopeScale	[headNum]	float16/bf16	ND	Mandatory. An empty tensor can be passed. This parameter is passed if cacheMode is set to 2. The data type is the same as that of input.

Output

Optional output tensors cannot be empty.

Module Function	Parameter	Dimension	Data Type	Format	Description
Output data	qOut0	If cacheMode is 0: [tokenNum,headNum, 576] If cacheMode is 1 or 2: [tokenNum,headNum, 512]	float16/bf16/int8	ND	Output tensor. Same data type as input. If cacheMode is 2, the data type is int8.
	kvCacheOut0	If cacheMode is 0: [blockNum, blockSize, 1, 576] If cacheMode is 1: [blockNum, blockSize, 1, 512] If cacheMode is 2: [blockNum, 1512/32, blockSize, 32] If cacheMode* is 3: [blockNum, 1*512/16, blockSize, 16]	float16/bf16/int8	ND/NZ	Output tensor. Same data type as input. If cacheMode is 2, the data type is int8 and the format is NZ. If cacheMode is 3, the format is NZ.
	qOut1	[tokenNum, headNum, 64]	float16/bf16	ND	This tensor is output if cacheMode is not set to 0. Same data type as input.
	kvCacheOut1	If cacheMode is 1: [blockNum, blockSize, 1, 64] If cacheMode is 2: [blockNum, 1*64/16, blockSize, 16]	float16/bf16	ND/NZ	This tensor is output if cacheMode is not set to 0. Same data type as input. If cacheMode is 2, the data format is NZ. If cacheMode is 3, the format is NZ.

Constraints

tokenNum ≤ 1024

blockSize ≤ 128 or blockSize = 256
If cacheMode is set to 2 or 3, blockSize is 128.
The value range of hiddenSize is [2048, 8192], and the value must be a multiple of 256.
If rmsNormQuant is disabled, hiddenSize can only be set to 6144, and cacheMode can only be set to 0 or 1.
If rmsNormQuant is disabled or hiddenSize is not set to 7168:
- 1 ≤ tokenNum ≤ 256
- Value range of headNum: [16, 32, 64, 128]

Parent topic: atb/infer_op_params.h