CANN/asc-devkit：SoftmaxFlash接口

SoftmaxFlash【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持Kirin X90支持Kirin 9030支持功能说明注意该接口后续即将废弃请使用精度和性能更好的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, inmaxNone, insumNone, updateNone): if update None: #基于last轴进行rowmax(按行取最大值)处理 x_max np.max(src, axis-1, keepdimsTrue) x_sub src - x_max x_exp np.exp(x_sub) #基于last轴进行rowsum(按行求和)处理 x_sum np.sum(x_exp, axis-1, keepdimsTrue) 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, keepdimsTrue) x_exp np.exp(src - x_max) x_sum np.sum(x_exp, axis-1, keepdimsTrue) 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 __aicore__ inline void SoftmaxFlash(const LocalTensorT dstTensor, const LocalTensorT sumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const LocalTensorT expMaxTensor, const LocalTensorT inSumTensor, const LocalTensorT inMaxTensor, const SoftMaxTiling tiling, bool isUpdate false, const SoftMaxShapeInfo softmaxShapeInfo {})template typename T, bool isReuseSource false, bool isBasicBlock false __aicore__ inline void SoftmaxFlash(const LocalTensorhalf dstTensor, const LocalTensorfloat sumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const LocalTensorhalf expMaxTensor, const LocalTensorfloat inSumTensor, const LocalTensorfloat inMaxTensor, const SoftMaxTiling tiling, bool isUpdate false, const SoftMaxShapeInfo softmaxShapeInfo {})通过sharedTmpBuffer入参传入临时空间template typename T, bool isReuseSource false, bool isBasicBlock false __aicore__ inline void SoftmaxFlash(const LocalTensorT dstTensor, const LocalTensorT sumTensor, const LocalTensorT maxTensor, const LocalTensorT srcTensor, const LocalTensorT expMaxTensor, const LocalTensorT inSumTensor, const LocalTensorT inMaxTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, bool isUpdate false, const SoftMaxShapeInfo softmaxShapeInfo {})template typename T, bool isReuseSource false, bool isBasicBlock false __aicore__ inline void SoftmaxFlash(const LocalTensorhalf dstTensor, const LocalTensorfloat sumTensor, const LocalTensorfloat maxTensor, const LocalTensorhalf srcTensor, const LocalTensorhalf expMaxTensor, const LocalTensorfloat inSumTensor, const LocalTensorfloat inMaxTensor, const LocalTensoruint8_t sharedTmpBuffer, const SoftMaxTiling tiling, bool isUpdate false, const SoftMaxShapeInfo softmaxShapeInfo {})由于该接口的内部实现中涉及复杂的计算需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。接口框架申请临时空间开发者无需申请但是需要预留临时空间的大小。通过sharedTmpBuffer入参传入使用该tensor作为临时空间进行处理接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间并在接口调用完成后复用该部分内存内存不会反复申请释放灵活性较高内存利用率也较高。接口框架申请的方式开发者需要预留临时空间通过sharedTmpBuffer传入的情况开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下通过SoftmaxFlash Tiling接口中提供的GetSoftMaxFlashMaxTmpSize/GetSoftMaxFlashMinTmpSize接口获取所需最大和最小临时空间大小最小空间可以保证功能正确最大空间用于提升性能。参数说明表1模板参数说明参数名描述T操作数的数据类型。支持的数据类型为half、float。isReuseSource该参数预留传入默认值false即可。isBasicBlocksrcTensor和dstTensor的shape信息和Tiling切分策略满足基本块要求的情况下可以设置为true开启该参数用于提升性能默认为false表示不开启。是否满足基本块的要求可以采用如下两种方式之一判断srcTensor和dstTensor的shape信息[m,n]需要满足如下条件尾轴长度n小于2048并且大于等于256/sizeof(T)即half场景下n最小为128float场景下n最小为64同时n是64的倍数非尾轴长度的乘积m为8的倍数。在Tiling实现中通过调用IsBasicBlockInSoftMax判断Tiling切分策略是否满足基本块的切分要求。表2接口参数说明参数名输入/输出描述dstTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。dstTensor的shape和源操作数srcTensor一致。sumTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于保存softmax计算过程中reducesum的结果。sumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock中的16个数均为相同的reducesum的值。非last轴的长度与dstTensor保持一致。maxTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于保存softmax计算过程中reducemax的结果。maxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如half数据类型下该datablock中的16个数均为相同的reducemax的值。非last轴的长度与dstTensor保持一致。srcTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。last轴长度需要32Byte对齐。expMaxTensor输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。expMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值。比如half数据类型下该datablock中的16个数均为相同的值。非last轴的长度与dstTensor保持一致。inSumTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。softmax计算所需要的sum值。inSumTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock中的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。inMaxTensor输入源操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。softmax计算所需要的max值。inMaxTensor的last轴长度固定为32Byte即一个datablock长度。该datablock中的所有数据为同一个值比如half数据类型下该datablock里的16个数均为相同的值。非last轴的长度需要与dstTensor保持一致。sharedTmpBuffer输入临时空间。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。接口内部复杂计算时用于存储中间变量由开发者提供。临时空间大小BufferSize的获取方式请参考SoftmaxFlash Tiling接口。tiling输入接口计算所需tiling信息Tiling信息的获取请参考SoftmaxFlash Tiling接口。isUpdate输入是否开启update算法。softmaxShapeInfo输入srcTensor的shape信息。SoftMaxShapeInfo类型具体定义如下方代码所示其中参数的含义为srcM非尾轴长度的乘积。srcK尾轴长度必须32Byte对齐。oriSrcM原始非尾轴长度的乘积。oriSrcK原始尾轴长度。struct SoftMaxShapeInfo { uint32_t srcM; uint32_t srcK; uint32_t oriSrcM; uint32_t oriSrcK; };返回值说明无约束说明srcTensor和dstTensor的空间可以复用maxTensor和inMaxTensor的空间可以复用sumTensor和inSumTensor的空间可以复用。sumTensor、maxTensor、expMaxTensor、inSumTensor、inMaxTensor的Tensor空间last轴长度必须固定32Byte。操作数地址对齐要求请参见通用地址对齐约束。不支持sharedTmpBuffer与源操作数和目的操作数地址重叠。调用示例本样例输入src的Shape大小为[80,144]输出Shape大小dst[80,144]输入inExpSumTensor[80,16]输入inMaxTensor[80,16]输出expMaxTensor[80,16]数据类型均为halfupdate为false。#include kernel_operator.h template typename T class KernelSoftmaxFlash { public: __aicore__ inline KernelSoftmaxFlash() {} __aicore__ inline void Init( GM_ADDR srcGm, GM_ADDR inMaxGm, GM_ADDR inSumGm, GM_ADDR dstGm, const SoftMaxTiling tilingData) { elementNumPerBlk 32 / sizeof(T); srcGlobal.SetGlobalBuffer((__gm__ T*)srcGm); maxGlobal.SetGlobalBuffer((__gm__ T*)inMaxGm); sumGlobal.SetGlobalBuffer((__gm__ T*)inSumGm); dstGlobal.SetGlobalBuffer((__gm__ T*)dstGm); pipe.InitBuffer(inQueueSrc, 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::LocalTensorT srcLocal inQueueSrc.AllocTensorT(); AscendC::LocalTensorT inSumLocal inSumQueue.AllocTensorT(); AscendC::LocalTensorT inMaxLocal inMaxQueue.AllocTensorT(); AscendC::DataCopy(srcLocal, srcGlobal, height * width); AscendC::DataCopy(inSumLocal, sumGlobal, height * elementNumPerBlk); AscendC::DataCopy(inMaxLocal, maxGlobal, height * elementNumPerBlk); inQueueSrc.EnQue(srcLocal); inSumQueue.EnQue(inSumLocal); inMaxQueue.EnQue(inMaxLocal); } __aicore__ inline void Compute() { AscendC::LocalTensorT srcLocal inQueueSrc.DeQueT(); AscendC::LocalTensorT dstLocal outQueueDst.AllocTensorT(); AscendC::LocalTensorT inMaxLocal inMaxQueue.AllocTensorT(); AscendC::LocalTensorT inSumLocal inSumQueue.AllocTensorT(); AscendC::LocalTensorT expMaxTensor expMaxQueue.AllocTensorT(); AscendC::SoftMaxShapeInfo srcShape {height, width, height, width}; AscendC::SoftmaxFlashT, false( srcLocal, inSumLocal, inMaxLocal, srcLocal, expMaxTensor, inSumLocal, inMaxLocal, tiling, false, srcShape); AscendC::DataCopy(dstLocal, srcLocal, height * width); outQueueDst.EnQueT(dstLocal); inMaxQueue.FreeTensor(inMaxLocal); inSumQueue.FreeTensor(inSumLocal); inQueueSrc.FreeTensor(srcLocal); expMaxQueue.FreeTensor(expMaxTensor); } __aicore__ inline void CopyOut() { AscendC::LocalTensorT dstLocal outQueueDst.DeQueT(); AscendC::DataCopy(dstGlobal, dstLocal, height * width); outQueueDst.FreeTensor(dstLocal); } private: AscendC::TPipe pipe; AscendC::TQueAscendC::TPosition::VECIN, 1 inQueueSrc; AscendC::TQueAscendC::TPosition::VECOUT, 1 outQueueDst; AscendC::TQueAscendC::TPosition::VECIN, 1 inMaxQueue; AscendC::TQueAscendC::TPosition::VECIN, 1 inSumQueue; AscendC::TQueAscendC::TPosition::VECIN, 1 expMaxQueue; AscendC::GlobalTensorT srcGlobal, dstGlobal; AscendC::GlobalTensorT maxGlobal, sumGlobal; uint32_t elementNumPerBlk 0; uint32_t width 144; uint32_t height 80; SoftMaxTiling tiling; }; extern C __global__ __aicore__ void softmax_flash_kernel_half( GM_ADDR srcGm, GM_ADDR inMaxGm, GM_ADDR inSumGm, GM_ADDR dstGm, GM_ADDR tiling) { GET_TILING_DATA(tilingData, tiling); KernelSoftmaxFlashhalf op; op.Init(srcGm, inMaxGm, inSumGm, dstGm, tilingData.softmaxTilingData); op.Process(); }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考