MlaPreprocessOperation

You can connect PyTorch by calling the C API 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	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, int64 if input is a float16 value, or float if input is a bf16 value.
bias0	Tensor	[2112]	int32	ND	-	Yes	Input tensor. An empty tensor can be passed. A 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, int64 if input is a float16 value, or float if input is a bf16 value.
bias1	Tensor	[headNum*192]	int32	ND	-	Yes	Input tensor. An empty tensor can be passed. A 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 0: [blockNum,blockSize, 1, 576] If cacheMode is 1: [blockNum,blockSize, 1, 512] If cacheMode is 2: [blockNum, headNum512/32, block_size, 32] If cacheMode* is 3: [blockNum, headNum*512/16, block_size, 16]	float16/bf16/int8	ND/NZ	-	Yes	Input tensor. Same data type as input. If cacheMode is 1, the tensor shape is split. If cacheMode is 2, the format is NZ and the data type is int8. If cacheMode is 3, the format is NZ.
kvCacheRope	Tensor	If cacheMode is 1: [blockNum,blockSize, 1, 64] If cacheMode is 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 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	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 0: [tokenNum,headNum, 576] If cacheMode is 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 2, the data type is int8.
kvCacheOut0	Tensor	If cacheMode is 0: [blockNum,blockSize, 1, 576] If cacheMode is 1: [blockNum,blockSize, 1, 512] If cacheMode is 2: [blockNum, headNum512/32, block_size, 32] If cacheMode* is 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 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. Same data type as input.
kvCacheOut1	Tensor	If cacheMode is 1: [blockNum,blockSize, 1, 64] If cacheMode is 2: [blockNum, headNum*64 / 16 , block_size, 16]	float16/bf16	ND/NZ	-	No	This tensor is output if cacheMode is not set to 0. Same data type as input. If cacheMode is set to 2, the data format is NZ. If cacheMode is set to 3, the data format is NZ.

Original API

See MlaPreprocessOperation.

Parent topic: atb/acl.h