RazorFusionAttentionOperation（代码开放）

产品支持情况

硬件型号	是否支持
Atlas A3 推理系列产品/Atlas A3 训练系列产品	√
Atlas A2 训练系列产品/Atlas 800I A2 推理产品	√
Atlas 训练系列产品	x
Atlas 推理系列产品	x
Atlas 200I/500 A2 推理产品	x

功能说明

支持混元大模型在多模态场景下的长序列优化，基于优化算法实现razorFusion特性。

图1 RazorFusionAttention参数示意图

blockSize=128（即图2）中每个方格的边长为128。

case1（图 a） (pre_tokens=5 * blockSize , next_tokens =4 * blockSize )
case2（图 b）(pre_tokens=5 * blockSize , next_tokens=6 * blockSize )

图2 pre_tokens和next_tokens参数示意图

定义

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};
};

参数列表

成员名称	类型	默认值	取值范围	是否必选	描述
headNum	int32_t	1	[1]	否	query头大小，当前只支持1。
kvHeadNum	int32_t	1	[1]	否	kv头数量，当前只支持1。
qkScale	float	1.0	-	否	算子tor值。
razorLen	int32_t	0	>=0	否	图片的长度。
preTokens	int32_t	0	>=0, 128对齐	否	用于稀疏计算，表示attention需要和前几个Token计算关联。
nextTokens	int32_t	0	>=0, 128对齐	否	用于稀疏计算，表示attention需要和前几个Token计算关联。
tileQ	int32_t	0	>=0	否	Q方向上图片的个数。
tileKv	int32_t	0	>=0	否	Kv方向图片的个数。
textQLen	int32_t	0	>=0	否	Q方向文本Token数量。
textKvLen	int32_t	0	>=0	否	Kv方向文本Token数量。

输入

参数	维度	数据类型	格式	描述
query	[b * s₁, h] / [b * s₁, 1, d]	float16/bf16	ND	query 矩阵。
key	[b * s₂, h] / [b * s₂, 1, d]	float16/bf16	ND	key 矩阵。
value	[b * s₂, h] / [b * s₂, 1, d]	float16/bf16	ND	value 矩阵。

输出

参数	维度	数据类型	格式	描述
output	[b * s1, h] / [b * s1, 1, d]	float16/bf16	ND	输出tensor。

约束说明

razorLen、tileQ、textQLen不能同时为0。
razorLen、tileKv、textKvLen不能同时为0。
s1（即qSeqLen）= razorLen * tileQ + textQLen，s2（即 kvSeqLen）= razorLen * tileKv + textKvLen，s1和s2不支持变长。
输入输出的数据类型需保持一致。
输入tensor最后一维h或d只支持128。

父主题： atb/infer_op_params.h