RelayAttentionOperation

Applicable Products

Hardware Model	Supported
Atlas 350 accelerator cards	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

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;
    uint8_t rsv[32] = {0};
};

Parameters

Member	Type	Default Value	Description
headNum	int32_t	0	Number of heads The value of headNum must be greater than 0.
qkScale	float	1.0	Tor value of the operator
kvHeadNum	int32_t	0	Number of kvheads. The value range is [0, 8]. The value of kvHeadNum can be either of the following: If the value is 0, kvHeadNum is equal to headNum. If the value is greater than 0, kvHeadNum is less than or equal to headNum and headNum is an integer multiple of kvHeadNum.
maskType	MaskType	MASK_TYPE_UNDEFINED	Mask type
rsv[32]	uint8_t	{0}	Reserved.

Input

Parameter	Dimension	Data Type	Format	CPU or NPU	Description
query	[B, qHiddenSize]	float16/bf16	ND	NPU	Query matrix.
key	[B, [S1, N, D]] [B, [S1, N*D]]	float16/bf16	ND	CPU	Non-shared key matrix, which is a TensorList. You can choose whether to combine axes.
value	[B, [S1, N, D]] [B, [S1, N*D]]	float16/bf16	ND	CPU	Non-shared value matrix, which is a TensorList. You can choose whether to combine axes.
keyShare	[BS, [S2, N, D]] [BS, [S2, N*D]]	float16/bf16	ND	CPU	Shared key matrix, which is a TensorList. You can choose whether to combine axes.
valueShare	[BS, [S2, N, D]] [BS, [S2, N*D]]	float16/bf16	ND	CPU	Shared value matrix, which is a TensorList. You can choose whether to combine axes.
attentionMask	-	float16/bf16	ND	NPU	Reserved input tensor.
seqLen	[B]	int32	ND	CPU	qseqlen list.
kvSeqLen	[B]	int32	ND	CPU	kvseqlen list.
kvShareMap	[B]	int32	ND	CPU	Array of the mapping between batches and shared groups. If there is no shared group, the value is -1.
kvShareLen	[BS]	int32	ND	CPU	Array of the actual length of the shared group.

Output

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

Description

Input key, value, keyShare, and valueShare are TensorLists. B and BS in these input tensors are level-2 pointer dimensions. Take key as an example. It is a list of size [B], which stores B pointers. Each pointer points to a tensor of size [S1, N, D] (or [S1, N × D]).
B indicates the batch. The value range is (0, 60]. BS indicates the number of shared groups. S1 indicates the length that is not shared. S2 indicates the shared length. N indicates kvHeadNum. D indicates headSize.
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.
Both key and value must be set to a combined axis or a non-combined axis at the same time. The same rule 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 can be set to MASK_TYPE_UNDEFINED only. Therefore, attentionMask is a reserved API and must be passed. Only the data type is restricted.

Parent topic: atb/infer_op_params.h