FastSoftMaxOperation

产品支持情况

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

功能描述

训练算子，与unpadOperation一起使用，将unpad处理后的Q矩阵和K矩阵相乘的结果做高性能的SoftMax处理。

SoftMax数学公式：

$\text{[math]}$

使用场景

图1 模型流程图

Multi-Head Attention计算过程如图1所示。经过unpad处理后，Q矩阵和K矩阵相乘的结果实际上由batchSize个大小为(seqLen[i], seqLen[i])的矩阵拼接而成（如图2所示），需要对这个结果的末轴做SoftMax操作。采用FastSoftMax算子，接收seqLen和headNum的信息可以对这种排列方式的数据实现SoftMax操作。

图2 输入数据的排列方式

定义

      
           struct FastSoftMaxParam {
    int32_t headNum = 0;
    std::vector<int32_t> qSeqLen;
    uint8_t rsv[8] = {0};
};

参数列表

成员名称	类型	默认值	描述
headNum	int32_t	0	Attention的head数量，需大于0。
qSeqLen	std::vector<int32_t>	-	每个batch的实际输入长度。元素个数为batchSize，最大不超过32。
rsv[8]	uint8_t	{0}	预留参数。