RelayAttentionOperation

Description

The core innovation of the RelayAttention algorithm is to group the matrix-vector multiplication corresponding to the system prompt into matrix-matrix multiplication, allowing the hidden state (key-value pair) of the system prompt to be read from the DRAM only once for a batch of input tokens, in this way, a large quantity of redundant memory accesses in existing causal attention calculation algorithms for processing system prompts are eliminated, and generation quality is maintained while efficiency is improved, and model retraining is not required.

Figure 1 RelayAttention computation principle

Operator Context

Definition

struct RelayAttentionParam {
    int32_t headNum = 0;
    float qkScale = 1;
    int32_t kvHeadNum = 0;
    enum MaskType : int {
        MASK_TYPE_UNDEFINED = 0,     
        MASK_TYPE_NORM,       
    };
    MaskType maskType = MASK_TYPE_UNDEFINED;
};

Parameters

Member	Type	Default Value	Description
headNum	int32_t	0	Number of heads
qkscale	float	1.0	Tor value of the operator
kvHeadNum	int32_t	0	Number of kvheads
maskType	MaskType	MASK_TYPE_UNDEFINED	Mask type

Input

Parameter	Dimension	Data Type	Format	CPU or NPU
query	[B, qHiddenSize]	float16/bf16	ND	NPU
key	[B, [S1, N ,D]] [B, [S1, N *D]]	float16/bf16	ND	CPU
value	[B, [S1, N ,D]] [B, [S1, N *D]]	float16/bf16	ND	CPU
keyShare	[BS, [S2, N ,D]] [BS, [S2, N *D]]	float16/bf16	ND	CPU
valueShare	[BS, [S2, N ,D]] [BS, [S2, N *D]]	float16/bf16	ND	CPU
attentionMask	-	float16/bf16	ND	NPU
seqLen	[B]	int32	ND	CPU
kvSeqLen	[B]	int32	ND	CPU
kvShareMap	[B]	int32	ND	CPU
kvShareLen	[BS]	int32	ND	CPU

Output

Parameter	Dimension	Data Type	Format	CPU or NPU
out	[nTokens, qHiddenSize]	float16/bf16	ND	NPU

Description

Only the Atlas A2 inference products is supported.
key, value, keyShare, and valueShare are TensorLists. B and BS in these input tensors are level-2 pointer dimensions.
B is batch, BS is the number of sharing groups, S1 is the length of non-sharing, S2 is the length of sharing, N is kvhead, and D is headdim.
CPU indicates that the input tensor is a host tensor. When an operator is passed, the tensor needs to be bound. NPU indicates that the input tensor is a device tensor.
For key and value, combining axis should be enabled or disabled at the same time. The same applies to keyShare and valueShare . That is, the dimensions of key and value are the same, and the dimensions of keyShare and valueShare are the same.
The input tensor must be of type float16 or bf16.
Currently, maskType supports only MASK_TYPE_UNDEFINED. Therefore, attentionMask is a reserved API and must be passed, but the content is not limited.

Parent topic: atb/infer_op_params.h