LaserAttentionOperation

产品支持情况

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

功能说明

训练场景下，使用LaserAttention算法实现self-attention（自注意力）的计算。

定义

struct LaserAttentionParam {
    int headNum = 0;
    std::string inputLayout = "BNSD";
    float scaleValue = 0.08838834764831843;
    float keepProb = 1.0f;
    int preTokens = 2147483647;
    int nextTokens = 1;
    int sparseMode = 0;
    int innerPrecise = 1;
    uint8_t rsv[8] = {0};
};

参数列表

成员名称	类型	默认值	取值范围	描述
headNum	int	0	(0, INT_MAX]	head个数，需大于0。
inputLayout	string	"BNSD"	"BNSD"	输入排布。
scaleValue	float	0.08838834764831843	(0, 1]	缩放系数。
keepProb	float	1.0f	目前暂未使用，配置值未生效，可配任意值	需要保留的比例，与dropMask相关。
preTokens	int	2147483647	[1, INT_MAX]，为256的整数倍	sparse纵向滑动窗口大小。与seqSize相关，约束参考约束说明。
nextTokens	int	1	目前暂未使用，配置值未生效，可配任意值	sparse横向滑动窗口大小。
sparseMode	int	0	目前暂未使用，配置值未生效，可配任意值	sparse模式，是否有下三角/梯形掩码。
innerPrecise	int	1	1	精度模式。1：高精度模式。
rsv[8]	uint8_t	{0}	-	预留参数。

输入

参数	维度	数据类型	格式	描述
query	[batch, qHeadNum, seqSize, q_head_size]	bf16	ND	注意力机制的Q矩阵。
key	[batch, kvHeadNum, kvSize, k_head_size]	bf16	ND	注意力机制的K矩阵。
value	[batch, kvHeadNum, kvSize, v_head_size	bf16	ND	注意力机制的V矩阵。
pseShift	-	-	-	目前为预留接口，输入空Tensor。
dropMask	-	-	-	目前为预留接口，输入空Tensor。
paddingMask	-	-	-	目前为预留接口，输入空Tensor。
attenMask	[seqSize, kvSize]	bf16	ND	attention范围。该接口可选，可正常输入，也可输入空Tensor不参与计算。
prefix	-	-	-	目前为预留接口，输入空Tensor。
actualSeqQLen	-	-	-	目前为预留接口，输入空Tensor。
actualSeqKVLen	-	-	-	目前为预留接口，输入空Tensor。

输出

参数	维度	数据类型	格式	描述
softmaxMax	[batch, qHeadNum, seqSize]	float	ND	输出tensor。
softmaxSum	[batch, qHeadNum, seqSize]	float	ND	输出tensor。
softmaxOut	[...]	bf16	ND	目前为预留接口，维度可任意配置。
attentionOut	[batch, qHeadNum, seqSize, v_head_size]	bf16	ND	输出tensor。

约束说明

qHeadNum值与参数headNum值相同。
qHeadNum值是kvHeadNum值的整数倍。
seqSize值和kvSize值均为256的整数倍。
构造query/key/value时，若值域配置在[-100, 100]以内，采取uniform均匀分布方式生成数据；否则，采取normal正态分布，均值在[-100, 100]内随机选取，标准差在[1, 25]内随机选取。
当attenMask不为空Tensor时。
- seqSize值和kvSize值相同。
- 为下三角形，当preTokens < seqSize时，算子内部当做梯形处理。
q_head_size、k_head_size和v_head_size的值有以下组合：
q_head_size

k_head_size

v_head_size

约束说明

128

128

128

seqSize值不小于参数preTokens值。

192

256

128

seqSize值和参数preTokens值相同。
该算子涉及但不支持确定性计算。
qHeadNum小于等于512。

q_head_size	k_head_size	v_head_size	约束说明
128	128	128	seqSize值不小于参数preTokens值。
192	256	128	seqSize值和参数preTokens值相同。

父主题： atb/train_op_params.h