SoftmaxFlash
功能说明
注意:该接口后续即将废弃,请使用精度和性能更好的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
函数原型
- 接口框架申请临时空间
1 2
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 = {})
1 2
template <typename T, bool isReuseSource = false, bool isBasicBlock = false> __aicore__ inline void SoftmaxFlash(const LocalTensor<half>& dstTensor, const LocalTensor<float>& sumTensor, const LocalTensor<float>& maxTensor, const LocalTensor<half>& srcTensor, const LocalTensor<half>& expMaxTensor, const LocalTensor<float>& inSumTensor, const LocalTensor<float>& inMaxTensor, const SoftMaxTiling& tiling, bool isUpdate = false, const SoftMaxShapeInfo& softmaxShapeInfo = {})
- 通过sharedTmpBuffer入参传入临时空间
1 2
template <typename T, bool isReuseSource = false, bool isBasicBlock = false> __aicore__ inline 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 LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, bool isUpdate = false, const SoftMaxShapeInfo& softmaxShapeInfo = {})
1 2
template <typename T, bool isReuseSource = false, bool isBasicBlock = false> __aicore__ inline void SoftmaxFlash(const LocalTensor<half>& dstTensor, const LocalTensor<float>& sumTensor, const LocalTensor<float>& maxTensor, const LocalTensor<half>& srcTensor, const LocalTensor<half>& expMaxTensor, const LocalTensor<float>& inSumTensor, const LocalTensor<float>& inMaxTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const SoftMaxTiling& tiling, bool isUpdate = false, const SoftMaxShapeInfo& softmaxShapeInfo = {})
由于该接口的内部实现中涉及复杂的计算,需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。
- 接口框架申请临时空间,开发者无需申请,但是需要预留临时空间的大小。
- 通过sharedTmpBuffer入参传入,使用该tensor作为临时空间进行处理,接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。
接口框架申请的方式,开发者需要预留临时空间;通过sharedTmpBuffer传入的情况,开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下:通过SoftmaxFlash Tiling接口中提供的GetSoftMaxFlashMaxTmpSize/GetSoftMaxFlashMinTmpSize接口获取所需最大和最小临时空间大小,最小空间可以保证功能正确,最大空间用于提升性能。
参数说明
参数名 |
描述 |
---|---|
T |
操作数的数据类型。 |
isReuseSource |
预留参数,暂未启用,为后续的功能扩展做保留,保持默认值即可。 |
isBasicBlock |
srcTensor和dstTensor的shape信息和Tiling切分策略满足基本块要求的情况下,可以使能该参数用于提升性能,默认不使能。基本块要求如下:
|
参数名 |
输入/输出 |
描述 |
---|---|---|
dstTensor |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float dstTensor的shape和源操作数srcTensor一致。 |
sumTensor |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float 用于保存softmax计算过程中reducesum的结果。
|
maxTensor |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float 用于保存softmax计算过程中reducemax的结果。
|
srcTensor |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float last轴长度需要32Byte对齐。 |
expMaxTensor |
输出 |
目的操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float
|
inSumTensor |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float softmax计算所需要的sum值。
|
inMaxTensor |
输入 |
源操作数。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品(Ascend 310P处理器)AI Core,支持的数据类型为:half/float softmax计算所需要的max值。
|
sharedTmpBuffer |
输入 |
临时空间。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 该操作数的数据类型固定uint8_t。 接口内部复杂计算时用于存储中间变量,由开发者提供。 临时空间大小BufferSize的获取方式请参考SoftmaxFlash Tiling接口。 |
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; // 原始尾轴长度 }; |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品(Ascend 310P处理器)AI Core
注意事项
- srcTensor和dstTensor的空间可以复用,maxTensor和inMaxTensor的空间可以复用,sumTensor和inSumTensor的空间可以复用。
- sumTensor、maxTensor、expMaxTensor、inSumTensor、inMaxTensor的Tensor空间,last轴长度必须固定32Byte。
- 操作数地址偏移对齐要求请参见通用约束。
调用示例
#include "kernel_operator.h" template <typename T> class KernelSoftmaxFlash { public: __aicore__ inline KernelSoftmaxFlash() {} __aicore__ inline void Init( __gm__ uint8_t *src1Gm, __gm__ uint8_t *src2Gm, __gm__ uint8_t *dstGm, const SoftMaxTiling &tilingData) { elementNumPerBlk = 32 / sizeof(T); src1Global.SetGlobalBuffer((__gm__ T *)src1Gm); src2Global.SetGlobalBuffer((__gm__ T *)src2Gm); dstGlobal.SetGlobalBuffer((__gm__ T *)dstGm); pipe.InitBuffer(inQueueSrc1, 1, height * width * sizeof(T)); pipe.InitBuffer(inQueueSrc2, 1, height * width * sizeof(T)); pipe.InitBuffer(outQueueDst, 1, height * width * sizeof(T)); pipe.InitBuffer(inMaxQueue, 1, height * elementNumPerBlk * sizeof(T)); pipe.InitBuffer(inSumQueue, 1, height * elementNumPerBlk * sizeof(T)); pipe.InitBuffer(expMaxQueue, 1, height * elementNumPerBlk * sizeof(T)); tiling = tilingData; } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> srcLocal1 = inQueueSrc1.AllocTensor<T>(); AscendC::LocalTensor<T> srcLocal2 = inQueueSrc2.AllocTensor<T>(); AscendC::DataCopy(srcLocal1, src1Global, height * width); AscendC::DataCopy(srcLocal2, src2Global, height * width); inQueueSrc1.EnQue(srcLocal1); inQueueSrc2.EnQue(srcLocal2); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> srcLocal1 = inQueueSrc1.DeQue<T>(); AscendC::LocalTensor<T> srcLocal2 = inQueueSrc2.DeQue<T>(); AscendC::LocalTensor<T> dstLocal = outQueueDst.AllocTensor<T>(); AscendC::LocalTensor<T> inmaxLocal = inMaxQueue.AllocTensor<T>(); AscendC::LocalTensor<T> insumLocal = inSumQueue.AllocTensor<T>(); AscendC::LocalTensor<T> expMaxTensor = expMaxQueue.AllocTensor<T>(); AscendC::SoftMaxShapeInfo srcShape = {height, width, height, width}; AscendC::SoftmaxFlash<T, false>(srcLocal2, insumLocal, inmaxLocal, srcLocal2, expMaxTensor, insumLocal, inmaxLocal, tiling, false, srcShape); AscendC::DataCopy(dstLocal, srcLocal2, height * width); outQueueDst.EnQue<T>(dstLocal); inMaxQueue.FreeTensor(inmaxLocal); inSumQueue.FreeTensor(insumLocal); inQueueSrc1.FreeTensor(srcLocal1); inQueueSrc2.FreeTensor(srcLocal2); expMaxQueue.FreeTensor(expMaxTensor); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> dstLocal = outQueueDst.DeQue<T>(); AscendC::DataCopy(dstGlobal, dstLocal, height * width); outQueueDst.FreeTensor(dstLocal); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc1; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueSrc2; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueueDst; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inMaxQueue; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inSumQueue; AscendC::TQue<AscendC::QuePosition::VECIN, 1> expMaxQueue; AscendC::GlobalTensor<T> src1Global, src2Global, dstGlobal; uint32_t elementNumPerBlk = 0; uint32_t width = 144; uint32_t height = 80; SoftMaxTiling tiling; }; extern "C" __global__ __aicore__ void softmax_flash_kernel_half( __gm__ uint8_t *src1Gm, __gm__ uint8_t *src2Gm, __gm__ uint8_t *dstGm, __gm__ uint8_t *tiling) { GET_TILING_DATA(tilingData, tiling); KernelSoftmaxFlash<half> op; op.Init(src1Gm, src2Gm, dstGm, tilingData.softmaxTilingData); op.Process(); }