RazorFusionAttentionOperation

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

Supports long sequence optimization of the Hunyuan-DiT in multi-modal scenarios and implements the razorFusion feature based on the optimization algorithm.

Figure 1 RazorFusionAttention parameter diagram

The side length of each grid in blockSize=128 (that is, Figure 2) is 128.

Case 1 (Figure a) (pre_tokens = 5 × blockSize, next_tokens = 4 × blockSize)
Case 2 (Figure b) (pre_tokens = 5 × blockSize, next_tokens = 6 × blockSize)

Figure 2 pre_tokens and next_tokens parameter diagram

Definition

struct RazorFusionAttentionParam {
    int32_t headNum = 1; 
    int32_t kvHeadNum = 1; 
    float qkScale = 1;
    int32_t razorLen = 0;
    int32_t preTokens = 0;
    int32_t nextTokens = 0;
    int32_t tileQ = 0;
    int32_t tileKv = 0;
    int32_t textQLen = 0;
    int32_t textKvLen = 0;
    uint8_t rsv[64] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
headNum	int32_t	1	[1]	No	Query header size. Currently, only 1 is supported.
kvHeadNum	int32_t	1	[1]	No	Number of KV heads. Currently, only 1 is supported.
qkScale	float	1.0	-	No	Tor value of the operator
razorLen	int32_t	0	>=0	No	Image length.
preTokens	int32_t	0	>= 0, 128-aligned	No	Used for sparse computation, indicating that the attention needs to be associated with the first several tokens.
nextTokens	int32_t	0	>= 0, 128-aligned	No	Used for sparse computation, indicating that the attention needs to be associated with the first several tokens.
tileQ	int32_t	0	>=0	No	Number of images in the Q direction.
tileKv	int32_t	0	>=0	No	Number of images in the Kv direction.
textQLen	int32_t	0	>=0	No	Number of text tokens in the Q direction.
textKvLen	int32_t	0	>=0	No	Number of text tokens in the Kv direction.

Input

Parameter	Dimension	Data Type	Format	Description
query	[b * s₁, h] / [b * s₁, 1, d]	float16/bf16	ND	Query matrix
key	[b * s₂, h] / [b * s₂, 1, d]	float16/bf16	ND	Key matrix
value	[b * s₂, h] / [b * s₂, 1, d]	float16/bf16	ND	Value matrix

Output

Parameter	Dimension	Data Type	Format	Description
output	[b * s1, h] / [b * s1, 1, d]	float16/bf16	ND	Output tensor.

Constraints

razorLen, tileQ, and textQLen cannot be set to 0 at the same time.
razorLen, tileKv, and textKvLen cannot be set to 0 at the same time.
s1 (that is, qSeqLen) = razorLen × tileQ + textQLen, s2 (that is, kvSeqLen) = razorLen × tileKv + textKvLen. s1 and s2 do not support variable lengths.
The input and output data types must be the same.
The last dimension h or d of the input tensor can only be 128.

Parent topic: atb/infer_op_params.h