注意：该接口后续即将废弃，请使用精度和性能更好的SoftmaxFlashV2接口。
SoftMax增强版本，除了可以对输入tensor做softmaxflash计算，还可以根据上一次softmax计算的sum和max来更新本次的softmax计算结果。last轴切轴的情况，每次计算的reduce结果并非是全轴的，需要根据上一次softmax计算的sum和max来更新本次的softmax计算结果，可以使用该增强接口。不支持NZ格式。
当前仅支持传入shape为ND格式，内部的reduce过程都是按last轴进行。不使能update时，该接口等同于SoftMax。
为方便理解，通过python脚本实现的方式，表达其计算公式如下，其中src、inmax、 insum、update为输入，dst、x_sum、x_max、exp_max为输出。
def softmax_flash(src, inmax=None, insum=None, update=None):
if update == None:
#基于last轴进行rowmax(按行取最大值)处理
x_max = np.max(src, axis=-1, keepdims=True)
x_sub = src - x_max
x_exp = np.exp(x_sub)
#基于last轴进行rowsum(按行求和)处理
x_sum = np.sum(x_exp, axis=-1, keepdims=True)
dst = x_exp / x_sum
exp_max = None
return dst, x_max, x_sum, exp_max
else:
#将inmax和src拼接后求rowmax
x_max = np.max(np.concatenate((inmax, src), axis=-1), axis=-1, keepdims=True)
x_exp = np.exp(src - x_max)
x_sum = np.sum(x_exp, axis=-1, keepdims=True)
exp_max = np.exp(inmax - x_max)
x_sum = exp_max * insum + x_sum
exp_max = exp_max * insum / x_sum
dst = x_exp / x_sum
return dst, x_max, x_sum, exp_max
template <typename T, bool isReuseSource = false, bool isBasicBlock = false>
void SoftmaxFlash(const LocalTensor<T>& dstTensor, const LocalTensor<T>& sumTensor,
const LocalTensor<T>& maxTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& expMaxTensor,
const LocalTensor<T>& inSumTensor, const LocalTensor<T>& inMaxTensor, const SoftMaxTiling& tiling,
bool isUpdate = false, const SoftMaxShapeInfo& softmaxShapeInfo = {})
参数名 |
输入/输出 |
描述 |
|---|---|---|
dstTensor |
输出 |
目的操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/float dstTensor的shape和源操作数srcTensor一致。 |
sumTensor |
输出 |
目的操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/float 用于保存softmax计算过程中reducesum的结果。
|
maxTensor |
输出 |
目的操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/float 用于保存softmax计算过程中reducemax的结果。
|
srcTensor |
输入 |
源操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/float last轴长度需要32Byte对齐。 |
expMaxTensor |
输出 |
目的操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/floatt
|
inSumTensor |
输入 |
源操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/floatt softmax计算所需要的sum值。
|
inMaxTensor |
输入 |
源操作数。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品，支持的数据类型为：half/float softmax计算所需要的max值。
|
tiling |
输入 |
接口计算所需tiling信息，Tiling信息的获取请参考SoftmaxFlash Tiling接口。 |
isUpdate |
输入 |
是否使能update算法。 |
softmaxShapeInfo |
输入 |
srcTensor的shape信息。SoftMaxShapeInfo类型，具体定义如下： struct SoftMaxShapeInfo {
uint32_t srcM; // 非尾轴长度的乘积
uint32_t srcK; // 尾轴长度，必须32Byte对齐
uint32_t oriSrcM; // 原始非尾轴长度的乘积
uint32_t oriSrcK; // 原始尾轴长度
};
|
isReuseSource |
输入 |
dstTensor是否复用srcTensor的空间。 |
isBasicBlock |
输入 |
srcTensor和dstTensor的shape信息和Tiling切分策略满足基本块要求的情况下，可以使能该参数用于提升性能，默认不使能。基本块要求如下：
|
无
Atlas A2训练系列产品