LaserAttentionGradOperation

功能

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

定义

struct LaserAttentionGradParam {
    int headNum = 0;
    std::string inputLayout = "BNSD";
    float scaleValue = 0.088;
    float keepProb = 1.0f;
    int preTokens = 2147483647;
    int nextTokens = 1;
    int sparseMode = 0;
    int innerPrecise = 1;
};

参数列表

成员名称	类型	默认值	描述
headNum	int	0	head个数。默认值为0。值大于0。
inputLayout	std::string	"BNSD"	输入排布。默认值为"BNSD"。仅支持配置为"BNSD"。
scaleValue	float	0.08838834764831843	缩放系数。默认值为0.088。目前暂未使用，配置值未生效。
keepProb	float	1.0f	需要保留的比例，与dropMask相关。默认值为1.0。目前暂未使用，配置值未生效。
preTokens	int	2147483647	sparse纵向滑动窗口大小。默认值为2147483647。取值范围为[1, seq_size]。取值为256的整数倍。
nextTokens	int	1	sparse横向滑动窗口大小。默认值为1。目前暂未使用，配置值未生效。
sparseMode	int	0	sparse模式，是否有下三角/梯形掩码。默认值为0。目前暂未使用，配置值未生效。
innerPrecise	int	1	精度模式。默认值为1，高精度模式。仅支持配置为1。

输入

参数	维度	数据类型	格式	描述
query	[batch, q_num_head, seq_size, head_dim]	bf16	ND	输入tensor。
key	[batch, kv_num_head, kv_size, head_dim]	bf16	ND	输入tensor。数据类型与query的数据类型相同。
value	[batch, kv_num_head, kv_size, head_dim]	bf16	ND	输入tensor。数据类型与query的数据类型相同。
attentionOutGrad	[batch, q_num_head, seq_size, head_dim]	bf16	ND	输入tensor。数据类型与query的数据类型相同。值域为[-0.5, 0.5]。
pseShift	[batch, q_num_head, seq_size, seq_size]/[batch, q_num_head, 1, kv_size]/[q_num_head, seq_size, seq_size]	bf16	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。数据类型与query的数据类型相同。
dropMask	[batch * q_num_head * seq_size * kv_size / 8]	uint8	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。
paddingMask	-	bf16	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。数据类型与query的数据类型相同。
attenMask	[seq_size, kv_size]	float16	ND	输入tensor，可选，不使用时传入空Tensor。attention范围。只支持下三角形。当preTokens < seq_size时，算子内部当做梯形处理。
softmaxMax	[batch, q_num_head, seq_size]	float	ND	输入tensor，为前向输出。
softmaxSum	[batch, q_num_head, seq_size]	float	ND	输入tensor，为前向输出。
softmaxIn	[...]	bf16	ND	输入tensor，为前向输出，可选，不使用时shape可任意配置，当前不支持使用。数据类型与query的数据类型相同。
attentionIn	[batch, q_num_head, seq_size, head_dim]	bf16	ND	输入tensor，为前向输出。数据类型与query的数据类型相同。
prefix	[batch]	int64	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。
actualSeqQLen	[batch]	int64	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。
actualSeqKVLen	[batch]	int64	ND	输入tensor，可选，不使用时传入空Tensor，当前不支持使用。

输出

参数	维度	数据类型	格式	描述
queryGrad	[batch, q_num_head, seq_size, head_dim]	bf16	ND	输出tensor。数据类型与query的数据类型相同。
keyGrad	[batch, kv_num_head, kv_size, head_dim]	bf16	ND	输出tensor。数据类型与query的数据类型相同。
valueGrad	[batch, kv_num_head, kv_size, head_dim]	bf16	ND	输出tensor。数据类型与query的数据类型相同。
dpse	[...]	bf16	ND	输出tensor，可选，不使用时shape可任意配置，当前不支持使用。数据类型与query的数据类型相同。

输入输出说明

可选：表示根据使用场景，可选择是否需要使用该Tensor。当不使用Tensor时：若为输入tensor，则传入空Tensor作为占位符，维度、数据类型等不受表中配置约束；若为输出tensor，shape可任意配置，维度、数据类型受表中配置约束。空Tensor为Tensor的默认初始状态，可参考Tensor接口说明，其维度数为0。
q_num_head的值与参数headNum的值相同。
seq_size的值大于等于参数preTokens的值。
q_num_head的值为kv_num_head的值的整数倍。
seq_size和kv_size的值是256的整数倍；当attenMask不为空Tensor时，seq_size和kv_size的值必须相同。
head_dim的值必须为128。
构造query/key/value时，若值域配置在[-100, 100]以内，采取uniform均匀分布方式生成数据；否则，采取normal正态分布，均值在[-100, 100]内随机选取，标准差在[1, 25]内随机选取。

规格约束

仅Atlas 800I A2 推理产品/Atlas A2 训练系列产品支持该算子。

该算子涉及但不支持确定性计算。

父主题： atb/train_op_params.h