RingMLAOperation

Applicable Products

Hardware Model	Supported or Not
Atlas 350 accelerator card	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

Based on the traditional MultiLatentAttention, the intermediate results lse and attention out output by the ring MLA operator are updated to global results, supporting a longer sequence length.

Definition

struct RingMLAParam {
    enum CalcType : int {
        CALC_TYPE_DEFAULT = 0,
        CALC_TYPE_FISRT_RING,  
        CALC_TYPE_MAX
    };
    enum KernelType : int {
        KERNELTYPE_DEFAULT = 0,   
        KERNELTYPE_HIGH_PRECISION 
    };
    enum MaskType : int {
        NO_MASK = 0,    
        MASK_TYPE_TRIU,
    };
    CalcType calcType = CalcType::CALC_TYPE_DEFAULT;
    int32_t headNum = 0;
    int32_t kvHeadNum = 0;
    float qkScale = 1;
    KernelType kernelType = KERNELTYPE_HIGH_PRECISION; 
    MaskType maskType = MASK_TYPE_TRIU;
    InputLayout inputLayout = TYPE_BSND;
    uint8_t rsv[64] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
calcType	CalcType	CALC_TYPE_DEFAULT	CALC_TYPE_DEFAULT CALC_TYPE_FISRT_RING	Yes	Computation type CALC_TYPE_DEFAULT: default. In the non-first and non-last card scenario, prevLse and prevOut are passed, and the softmaxLse output is generated. CALC_TYPE_FISRT_RING: In the first-card scenario, no prevLse or prevOut is passed, and the softmaxLse output is generated.
headNum	int32_t	0	Greater than 0	Yes	Query head size, which must be greater than 0.
kvHeadNum	int32_t	0	Greater than or equal to 0	Yes	Number of KV heads. The value is determined based on the actual model. If kvHeadNum is 0, the values of kHeadNum of key and vHeadNum of value are the same as the value of headNum of query. If kvHeadNum is not set to 0, the values of kHeadNum of key and vHeadNum of value are the same as the value of kvHeadNum, and the value of headNum can be exactly divided by the value of kvHeadNum.
qkScale	float	1	-	Yes	Tor value of the operator, which is post-multiplied by Q*K^T.
kernelType	KernelType	KERNELTYPE_HIGH_PRECISION	KERNELTYPE_HIGH_PRECISION	Yes	Kernel precision type KERNELTYPE_HIGH_PRECISION: The input/output tensor uses the float16/bf16 type, and softmax uses the float type.
maskType	MaskType	MASK_TYPE_TRIU	NO_MASK MASK_TYPE_TRIU	Yes	Mask type NO_MASK: No mask is used. MASK_TYPE_TRIU: upper triangular mask. This is the default value.
inputLayout	InputLayout	TYPE_BSND	TYPE_BSND	Yes	Data layout format. The default data layout format is BSND.
rsv[64]	uint8_t	{0}	[0]	No	Reserved

Input

Parameter	Dimension	Data Type	Format	cpu/npu	Description	Application Scenarios
query	[qNTokens, headNum, 128]	float16/bf16	ND	npu	Query matrix without positional encoding.	Basic scenario
queryRope	[qNTokens, headNum, 64]	float16/bf16	ND	npu	Query component with rotational positional encoding.	Basic scenario
key	[kvNTokens, kvHeadNum, 128]	float16/bf16	ND	npu	Key matrix without positional encoding.	Basic scenario
keyRope	[kvNTokens, kvHeadNum, 64]	float16/bf16	ND	npu	Key with rotational positional encoding.	Basic scenario
value	[kvNTokens, kvHeadNum, 128]	float16/bf16	ND	npu	Value matrix	Basic scenario
mask	[512, 512] [batch, 512, 512]	float16/bf16	ND	npu	Mask with the same batch. Mask with different batches.	Basic scenario
seqLen	[batch]/[2, batch]	int32/uint32	ND	cpu	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, key, and value. If the shape is [2, batch], seqLen[0] indicates the sequence length of query, and seqLen[1] indicates the sequence length of key and value.	Basic scenario
prevOut	[qNTokens, headNum, 128]	float16/bf16	ND	npu	Previous output.	Non-first-card scenario
prevLse	[headNum, qNTokens]	float	ND	npu	Result of the previous QK^T × tor. softmax, exp, sum, and log are calculated in sequence.	Non-first-card scenario

Output

Parameter	Dimension	Data Type	Format	cpu/npu	Description	Application Scenarios
output	[qNTokens, headNum, headSizeV]	float16/bf16	ND	npu	It has the same address as prevOut.	Basic scenario
softmaxLse	[headNum, qNTokens]	float	ND	npu	It has the same address as prevLse.	Basic scenario

Functions

First-card scenario
- How to enable: calcType = CALC_TYPE_FISRT_RING
- Difference: No prevLse or prevOut is passed, and the softmaxLse output is generated.
Non-first-card scenario
- How to enable: calcType = CALC_TYPE_DEFAULT
- Difference: prevLse and prevOut are passed, and the softmaxLse output is generated.

Constraints

mask is used only when maskType is MASK_TYPE_TRIU.
inputLayout supports only TYPE_BSND.
Restrictions on 2D seqLen:
- qSeqLen is seqLen[0].
- kvSeqLen is seqLen[1].
- For each subscript i, qSeqLen[i] cannot be 0.
- If mask is used, for each subscript i, kvSeqLen[i] can be 0. The output result is the value of prevOut corresponding to seqlen in each batch. Note that kvSeqLen[0] and kvSeqLen[batch - 1] cannot be 0.
- If calcType is set to CALC_TYPE_FISRT_RING, for each subscript i, kvSeqLen[i] must be greater than or equal to qSeqLen[i].

Parent topic: atb/infer_op_params.h