RingMLAOperation

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

Definition

atb::Status AtbRingMLAGetWorkspaceSize(const aclTensor *querySplit1, const aclTensor *querySplit2,
                                       const aclTensor *keySplit1, const aclTensor *keySplit2, const aclTensor *value,
                                       const aclTensor *mask, const aclTensor *seqLen, const aclTensor *prevOut,
                                       const aclTensor *prevLse, int32_t headNum, int32_t kvHeadNum, float qkScale,
                                       int kernelType, int maskType, int inputLayout, int calcType, aclTensor *output,
                                       aclTensor *softmaxLse, uint64_t *workspaceSize, atb::Operation **op,
                                       atb::Context *context);
atb::Status AtbRingMLA(void *workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);

AtbRingMLAGetWorkspaceSize Members

Parameter	Scalar/Tensor	Dimension	Data Type	Format	Default Value	Mandatory or Not	Description
querySplit1	Tensor	[qNTokens, headNum, 128]	float16/bf16	ND	-	Yes	Input tensor, query matrix without positional encoding.
querySplit2	Tensor	[qNTokens, headNum, 64]	float16/bf16	ND	-	Yes	Input tensor, query component with rotational positional encoding.
keySplit1	Tensor	[kvNTokens, kvHeadNum, 128]	float16/bf16	ND	-	Yes	Input tensor, key matrix without positional encoding.
keySplit2	Tensor	[kvNTokens, kvHeadNum, 64]	float16/bf16	ND	-	Yes	Input tensor, key with rotational positional encoding.
value	Tensor	[kvNTokens, kvHeadNum, 128]	float16/bf16	ND	-	Yes	Input tensor, value matrix.
mask	Tensor	[512, 512]	float16/bf16	ND	-	Yes	Input tensor, mask.
seqLen	Tensor	[batch]/[2, batch]	int32/uint32	ND	-	Yes	Input tensor, sequence length. If the shape is [batch], it indicates the sequence length of each batch. The sequence length of each batch is the same for query, cacheK, and cacheV. If the shape is [2, batch], seqlen[0] indicates the sequence length of query, and seqlen[1] indicates the sequence length of cacheK and cacheV.
prevOut	Tensor	[qNTokens, headNum, headSizeV]	float16/bf16	ND	-	No	Input tensor, previous output.
prevLse	Tensor	[headNum, qNTokens]	float16	ND	-	No	Input tensor, result of the previous QK^T × tor. softmax, exp, sum, and log are calculated in sequence.
headNum	Scalar	-	int32_t	-	0	Yes	Size of the query head.
kvHeadNum	Scalar	-	int32_t	-	0	Yes	Number of KV heads. The value is determined based on the actual model. When kvHeadNum is 0, k_head_num of keyCache, v_head_num of valueCache, and num_heads of query are the same, that is, the value of num_heads. When kvHeadNum is not 0, k_head_num of keyCache and v_head_num of valueCache are the same as kvHeadNum.
qkScale	Scalar	-	float	-	1.0	Yes	Tor value of the operator, which is post-multiplied by Q*K^T.
kernelType	Scalar	-	int	-	1	Yes	Kernel precision type. Only 1 indicating KERNELTYPE_HIGH_PRECISION is supported.
maskType	Scalar	-	int	-	1	Yes	Mask type. 0: NO_MASK 1: MASK_TYPE_TRIU
inputLayout	Scalar	-	int	-	0	Yes	Data layout format. Only 0 (TYPE_BSND) is supported.
calcType	Scalar	-	int	-	0	Yes	Computation type. CALC_TYPE_DEFAULT = 0: default. In the non-first and non-last card scenario, prev_lse and prev_o are passed, and the softmaxLse output is generated. Currently, only the default value is supported. CALC_TYPE_FISRT_RING: In the first-card scenario, no prev_lse or prev_o is passed, and the softmaxLse output is generated.
output	Tensor	[qNTokens, headNum, headSizeV]	float16/bf16	ND	-	Yes	Output tensor.
softmaxLse	Tensor	[headNum, qNTokens]	float	ND	-	Yes	softmaxLse output.

Original APIs

For details, see RingMLAOperation.

Parent topic: atb/acl.h