MLA

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

Definition

atb::Status AtbMLAGetWorkspaceSize(const aclTensor *qNope, const aclTensor *qRope, const aclTensor *ctKV,
                                   const aclTensor *kRope, const aclTensor *blockTables, const aclTensor *contextLens,
                                   const aclTensor *mask, const aclTensor *qSeqLen, const aclTensor *qkDescale,
                                   const aclTensor *pvDescale, int32_t headNum, float qkScale, int32_t kvHeadNum,
                                   int maskType, int calcType, uint8_t cacheMode, aclTensor *attenOut, aclTensor *ise,
                                   uint64_t *workspaceSize, atb::Operation **op, atb::Context *context);
atb::Status AtbMLA(void* workspace, uint64_t workspaceSize, atb::Operation *op, atb::Context *context);

AtbMLAGetWorkspaceSize Members

Parameter	Scalar/Tensor	Dimension	Data Type	Format	Default Value	Mandatory or Not	Description
qNope	Tensor	[num_tokens, num_heads, 512]	float16/bf16/int8	ND	-	Yes	Query without positional encoding.
qRope	Tensor	[num_tokens, num_heads, 64]	float16/bf16	ND	-	Yes	Query with rotational positional encoding.
ctKV	Tensor	[num_blocks , block_size, kv_heads, 512] If cacheMode is set to 2: [blockNum, kv_heads512/32,block_size, 32] If cacheMode* is set to 3: [blockNum, kv_heads*512/16,block_size, 16]	float16/bf16/int8	ND/NZ	-	Yes	ctkv without positional encoding. If cacheMode is set to 2, the value is of the int8 data type and the format is NZ. If cacheMode is set to 3: the value is of the float16/bf16 data type and the format is NZ.
kRope	Tensor	[num_blocks , block_size, kv_heads, 64] If cacheMode is set to 2 or 3: [blockNum, kv_heads*64 / 16 ,block_size, 16]	float16/bf16	ND/NZ	-	Yes	Key with rotational positional encoding. If cacheMode is set to 2 or 3, NZ is used.
blockTables	Tensor	[batch, max_num_blocks_per_query]	int32	ND	-	Yes	kvcache block mapping table of each query.
contextLens	Tensor	[batch]	int32	ND	-	Yes	Context length (kseqlen) corresponding to each query. Tensor on the host side.
mask	Tensor	MASK_TYPE_SPEC: [num_tokens (combined axis), max_seq_len] MASK_TYPE_MASK_FREE: [125 + 2 * qseqlen, 128]	float16/bf16	ND	-	No	Attention mask. If maskType is set to 0, nullptr is passed.
qseqlen	Tensor	[batch]	int32	ND	-	No	seqLen corresponding to each batch. This parameter is passed when calcType is set to 1.
qkDescale	Tensor	[num_heads]	float	ND	-	No	This parameter is passed when cacheMode is set to 2.
pvDescale	Tensor	[num_heads]	float	ND	-	No	This parameter is passed when cacheMode is set to 2.
headNum	Scalar	-	int32_t	-	0	Yes	Number of query headers.
qkScale	Scalar	-	float	-	1.0	Yes	Scaling coefficient post-multiplied by Q*K^T.
kvHeadNum	Scalar	-	int32_t	-	0	Yes	Number of KV heads.
maskType	Scalar	-	int	-	0	No	Mask type. UNDEFINED (default): all-0 mask. MASK_TYPE_SPEC: mask when qseqlen > 1. MASK_TYPE_MASK_FREE: maskfree function.
calcType	Scalar	-	int	-	0	No	Computation type. CALC_TYPE_UNDEFINED = 0: default value. CALC_TYPE_SPEC: qseqlen greater than 1 can be passed. CALC_TYPE_RING: ringAttention. CALC_TYPE_PREFILL: full prefill.
cacheMode	Scalar	-	int	-	0	Yes	KVCACHE = 0: concatenated cache. KROPE_CTKV: separated cache. INT8_NZCACHE: high-performance separated cache. NZCACHE: non-quantized NZcache.
attenOut	Tensor	[num_tokens, num_heads, 512]	float16/bf16	ND	-	Yes	Attention output.
lse	Tensor	[num_tokens, num_heads, 1]	float16/bf16	ND	-	No	LSE output. If calcType is not CALC_TYPE_RING, nullptr can be passed.

Original APIs

For details, see Input and Output.

Parent topic: MultiLatentAttentionOperation