SoftmaxV2 ARA FullLoad 模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 模板用途针对 3D[A1, R, A0]布局、R×A0_tile 可载入 UB 的 Softmax 场景。沿 A0 切 tile跨 R 逐列 VF 计算根据 R 大小分支选择不同的归约路径R≤2、R≤4、R≤8 直接累加R8 使用 BinaryAddVF 二分折叠。2. 输入输出布局和 Softmax 轴输入x[A1, R, A0]行优先Softmax 沿 R 轴计算。输出y[A1, R, A0]与输入同 shape、同 dtype。数据在 UB 中布局为(R, tileA0Len)沿 R 逐行处理A0 方向向量化。3. 适用 shape、dtype 和 UB 条件条件要求输入维数3DA0 1R 大小R×tileA0Len 可载入 UBdtypeFP32、FP16、BF16UB 约束tileA0Len × (R1) × (dtypeSize×2 4) tileA0Len × 4≤ UB4. 类名、构造函数和 Init 接口namespace SoftmaxV2Ops; template typename T1, typename T2 class SoftmaxV2ARA { public: __aicore__ inline SoftmaxV2ARA(); __aicore__ inline SoftmaxV2ARA(const SoftmaxV2ARATilingData* tilingDataIn); __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, TPipe* pipeIn); __aicore__ inline void Process(); };构造函数无参默认构造或接收 TilingData 指针。Init接收 x/y GM 地址和TPipe*初始化 GlobalTensor 和 UB BufferxQueue/yQueue 双缓冲 xTmpBuf xReduceBuf。Process按 tile 遍历所有 (A1, A0) 组合逐 tile 计算 Softmax。注意ARA 系模板的构造/Init 方式与 AR 系不同——ARA 在 Init 中接收TPipe*而 AR 在构造函数中接收。5. 对应的 TilingData 和字段说明对应SoftmaxV2ARATilingData定义在 softmax_v2_tiling_data.h。主要字段字段含义单位totalA1Len / totalRLen / totalA0Len各维总长度元素数totalTiles总 tile 数 A1 × a0OutertiletilesPerCore每核 tile 数tileusedCoreNums实际使用核数核a0OuterA0 方向 tile 总数tiletileA0Len / tileA0Tail主块/尾块 A0 tile 长度元素数binaryAddKBinaryAdd 外层折叠层数层binaryAddLast最后是否需额外 2→1 折叠标志binaryAddInnerLoopBinaryAdd 每层内循环次数次remainderLoopCountR 对 8 取余后的折叠循环数次quotientLoopCountR 整除 8 后的剩余折叠循环数次remainderTailOffset0~7余数尾块各行偏移越界指向 validNumInXUb 做零填充元素数6. Host tiling 参数计算方法使用 softmax_v2_tiling.h 中的TilingAraFullLoadsoftmax_tiling::CaseShape shape{a1, r, a0, dtypeCode}; softmax_tiling::PlatformParam plat{ubSize, numBlocks}; SoftmaxV2ARATilingData td; int64_t blockDim softmax_tiling::TilingAraFullLoad(shape, plat, td);关键计算a0TileBase FP32 ? 64 : 128tileA0Len a0Inner × a0TileBase受 UB 容量和核数约束当 R ≤ 8 时binaryAdd* 字段不使用当 R 8 时binaryAddQuotient 最大 2^k ≤ R-1按 8 行分组做 BinaryAddVF7. CopyIn、归约、归一化和 CopyOut 流程Process: for 每个 tile (curA1Idx, curA0Idx): 1. CopyInX: DataCopyPad 批量载入 (curTileRLen × curTileA0Len) 块 2. Compute: a. VFShiftVector: 逐 A0 列求 max → sub max → exp → 写入 xTmpLocal b. VFReduceSum: 按 R 大小分支求 sum c. VFCalculateOutput: 逐列 div sum → Cast 回原精度 3. CopyOutY: DataCopyPad 写回 GM8. VF、BinaryAddVF 优化R 大小分支归约R 范围归约路径说明R ≤ 2SumRLessThan2逐行 VFAdd累加R ≤ 4SumRLessThan42 行配对 尾行处理TwoRowAddWithTailR ≤ 8SumRLessThan84 行配对 尾行处理R 8SumRMoreThan8BinaryAddVF8 行分组 二分折叠BinaryAddVF 二分折叠R 8行分组将 R 行按 8 行一组分组每组内用TwoRowAdd做 4→2→1 折叠。余数处理不满 8 行的余数部分用remainderTailOffset0~7精确控制每行偏移越界行指向validNumInXUb零填充区域。二分树BinaryAddVF按binaryAddK层 ×binaryAddInnerLoop次每次 4→1 折叠最后若binaryAddLast1额外做 2→1。TwoRowAddWithTail通用工具函数处理 2 行配对 2 个尾行偏移。9. 主块、尾块、对齐和数值精度处理主块/尾块tileA0Len为主块 A0 长度tileA0Tail为最后一个 A0 tile 的实际长度。尾块通过copyInParams.dstStride和copyInParams.srcStride处理非对齐。对齐A0 按a0TileBaseFP3264/FP16128对齐dstStride按BLOCK_SIZE32B对齐。零填充R 8 时在xTmpLocal[validNumInXUb]位置写入零值余数尾块越界行指向该区域等效于零填充。数值精度FP16/BF16 通过LoadTensorForDtypeT1DIST_UNPACK_B16Cast升至 FP32输出通过CastDIST_PACK_B32降回。10. 独立 Kernel 入口示例#include softmax_v2_tiling_data.h #include dav310/softmax_v2_base.h #include dav310/softmax_v2_ara_full_load.h using namespace SoftmaxV2Ops; extern C __global__ __aicore__ __vector__ void softmax_ara_full_load_fp32(GM_ADDR x, GM_ADDR y, SoftmaxV2ARATilingData tiling) { TPipe pipe; SoftmaxV2ARAfloat, float op(tiling); op.Init(x, y, pipe); op.Process(); }11. 不适用场景和可选替代模板不适用场景推荐替代R×tileA0 超 UBARA RecomputeR 切 bin 三阶段重读A0 12DAR FullLoad 或 AR SmallR需要在线更新避免 3 次读入ARA Online2 次读入【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考