MlaPreprocessOperation

The PyTorch can be directly connected through the C interface to replace affinity operators on the entire network.

Definition

atb::Status AtbMLAPreprocessGetWorkspaceSize(
            const aclTensor *input, const aclTensor *gamma0, const aclTensor *beta0, const aclTensor *quantScale0,
            const aclTensor *quantOffset0, const aclTensor *wdqkv, const aclTensor *deScale0, const aclTensor *bias0,
            const aclTensor *gamma1, const aclTensor *beta1, const aclTensor *quantScale1, const aclTensor *quantOffset1,
            const aclTensor *wuq, const aclTensor *deScale1, const aclTensor *bias1, const aclTensor *gamma2,
            const aclTensor *cos, const aclTensor *sin, const aclTensor *wuk, const aclTensor *kvCache,
            const aclTensor *kvCacheRope, const aclTensor *slotmapping, const aclTensor *ctkvScale, const aclTensor *qNopeScale,
            uint32_t wdqDim, uint32_t qRopeDim, uint32_t kRopeDim, float epsilon, uint32_t qRotaryCoeff, uint32_t kRotaryCoeff,
            bool transposeWdq, bool transposeWuq, bool transposeWuk, uint8_t cacheMode, uint16_t quantMode, aclTensor *qOut0,
            aclTensor *kvCacheOut0, aclTensor *qOut1, aclTensor *kvCacheOut1, uint64_t *workspaceSize, atb::Operation **op,
            atb::Context *context);
atb::Status AtbMLAPreprocess(void *workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);

AtbMLAPreprocessGetWorkspaceSize Members

Parameter	Scalar/Tensor	Dimension	Data Type	Format	Default Value	Mandatory or Not	Description
input	Tensor	[tokenNum, 7168]	float16/bf16	ND	-	Yes	Input tensor.
gamma0	Tensor	[7168]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
beta0	Tensor	[7168]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
quantScale0	Tensor	[1]	float16/bf16	ND	-	Yes	Input tensor. 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	Tensor	[1]	int8	ND	-	Yes	Input tensor. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0.
wdqkv	Tensor	[1,224,2112,32]	int8	NZ	-	Yes	Input tensor.
deScale0	Tensor	[2112]	int64/float	ND	-	Yes	Input tensor. 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	Tensor	[2112]	int32	ND	-	Yes	Input tensor. An empty tensor can be passed. nullptr is passed when quantMode is set to 1 or 3.
gamma1	Tensor	[1536]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
beta1	Tensor	[1536]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
quantScale1	Tensor	[1]	float16/bf16	ND	-	Yes	Input tensor. 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	Tensor	[1]	int8	ND	-	Yes	Input tensor. An empty tensor can be passed. This parameter is passed only if quantMode is set to 0.
wuq	Tensor	[1,48,headNum*192,32]	int8	NZ	-	Yes	Input tensor.
deScale1	Tensor	[headNum*192]	int64/float	ND	-	Yes	Input tensor. 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	Tensor	[headNum*192]	int32	ND	-	Yes	Input tensor. An empty tensor can be passed. nullptr is passed when quantMode is set to 1 or 3.
gamma2	Tensor	[512]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
cos	Tensor	[tokenNum,64]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
sin	Tensor	[tokenNum,64]	float16/bf16	ND	-	Yes	Input tensor. The data type is the same as that of input.
wuk	Tensor	ND: [headNum, 128, 512] NZ: [headNum, 32, 128, 16]	float16/bf16	ND/NZ	-	Yes	Input tensor. The data type is the same as that of input.
kvCache	Tensor	If cacheMode is set to 0: [blockNum,blockSize,1,576] If cacheMode is set to 1: [blockNum,blockSize,1,512] If cacheMode is set to 2: [blockNum, headNum512/32,block_size, 32] If cacheMode* is set to 3: [blockNum, headNum*512/16,block_size, 16]	float16/bf16/int8	ND/NZ	-	Yes	Input tensor. The data type is the same as that of input. If cacheMode is set to 1, the tensor shape is split. If cacheMode is set to 2, the format is NZ and the data type is int8. If cacheMode is set to 3, the format is NZ.
kvCacheRope	Tensor	If cacheMode is set to 1: [blockNum,blockSize,1,64] If cacheMode is set to 2 or 3: [blockNum, headNum*64 / 16 ,block_size, 16]	float16/bf16	ND/NZ	-	Yes	Input tensor. 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 set to 2 or 3, the format is NZ.
slotmapping	Tensor	[tokenNum]	int32	ND	-	Yes	Input tensor.
ctkvScale	Tensor	[1]	float16/bf16	ND	-	Yes	Input tensor. 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.
qNopeScale	Tensor	[headNum]	float16/bf16	ND	-	Yes	Input tensor. 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.
wdqDim	Scalar	-	uint32_t	-	0	Yes	Split dimension size after matmul.
qRopeDim	Scalar	-	uint32_t	-	0	Yes	Dimension size of q passed to RoPE.
kRopeDim	Scalar	-	uint32_t	-	0	Yes	Dimension size of k passed to RoPE.
epsilon	Scalar	-	float	-	1e-5	Yes	It is added to the denominator to prevent division by 0.
qRotaryCoeff	Scalar	-	int32_t	-	2	Yes	q rotation coefficient.
kRotaryCoeff	Scalar	-	int32_t	-	2	Yes	k rotation coefficient.
transposeWdq	Scalar	-	bool	-	true	Yes	Whether to transpose wdq.
transposeWuq	Scalar	-	bool	-	true	Yes	Whether to transpose wuq.
transposeWuk	Scalar	-	bool	-	true	Yes	Whether to transpose wuk.
cacheMode	Scalar	-	uint8_t	-	0	Yes	0: KVCACHE 1: KROPE_CTKV 2: INT8_NZCACHE 3: NZCACHE
quantMode	Scalar	-	uint16_t	-	0	Yes	0: PER_TENSOR_QUANT_ASYMM 1: PER_TOKEN_QUANT_SYMM 2: PER_TOKEN_QUANT_ASYMM 3: UNQUANT
qOut0	Tensor	If cacheMode is set to 0: [tokenNum,headNum,576] If cacheMode is set to 1 or 2: [tokenNum,headNum,512]	float16/bf16/int8	ND	-	Yes	Output tensor. The data type is the same as that of input. If cacheMode is set to 2, the data type is int8.
kvCacheOut0	Tensor	If cacheMode is set to 0: [blockNum,blockSize,1,576] If cacheMode is set to 1: [blockNum,blockSize,1,512] If cacheMode is set to 2: [blockNum, headNum512/32,block_size, 32] If cacheMode* is set to 3: [blockNum, headNum*512/16,block_size, 16]	float16/bf16/int8	ND/NZ	-	Yes	Output tensor. The data type is the same as that of input. If cacheMode is set to 2, the data type is int8 and the format is NZ. If cacheMode is set to 3, the format is NZ.
qOut1	Tensor	[tokenNum,headNum,64]	float16/bf16	ND	-	No	This tensor is output if cacheMode is not set to 0. The data type is the same as that of input.
kvCacheOut1	Tensor	If cacheMode is set to 1: [blockNum,blockSize,1,64] If cacheMode is set to 2: [blockNum, headNum*64 / 16 ,block_size, 16]	float16/bf16	ND/NZ	-	No	This tensor is output if cacheMode is not set to 0. The data type is the same as that of input. If cacheMode is set to 2, the data format is NZ. If cacheMode is set to 3, the data format is NZ.

Original APIs

For details, see MlaPreprocessOperation.

Parent topic: atb/acl.h