SoftmaxV2 ARA Recompute 模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 模板用途针对 3D[A1, R, A0]布局、R 超出 UB 容量的 Softmax 场景。将 R 切分为多个 bin通过三阶段重读 GM求 max → 求 sum → 输出完成 Softmax。使用NlastReduceSum跨 bin 二分折叠和UpdateCache累加树高效合并局部 sum。2. 输入输出布局和 Softmax 轴输入x[A1, R, A0]行优先Softmax 沿 R 轴计算。输出y[A1, R, A0]与输入同 shape、同 dtype。UB 内布局为(binAddRFactor, tileA0Len)沿 Rbin逐块处理。3. 适用 shape、dtype 和 UB 条件条件要求输入维数3DA0 1R 大小R 超 UB 单次载入容量dtypeFP32、FP16、BF16UB 约束双缓冲 xQueue/yQueue xMax/xSum cache/temp/reduceSumTemp buffer4. 类名、构造函数和 Init 接口namespace SoftmaxV2Ops; template typename T1, typename T2 class SoftmaxV2ARARecompute : public SoftmaxV2OpsBase { public: __aicore__ inline SoftmaxV2ARARecompute(); __aicore__ inline SoftmaxV2ARARecompute(const SoftmaxV2ARARecomputeTilingData* tilingDataIn); __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, TPipe* pipeIn); __aicore__ inline void Process(); };构造函数无参默认或接收 TilingData 指针。Init初始化 GM/UB buffer包括 xQueue/yQueue双缓冲、xMaxBuf/xSumBuf、cacheBuffer/tempBuffer/reduceSumTempBuffer。Process按 tile 遍历每 tile 执行三阶段计算。5. 对应的 TilingData 和字段说明对应SoftmaxV2ARARecomputeTilingData定义在 softmax_v2_tiling_data.h。字段含义单位totalRLen / totalA0LenR/A0 维总长度元素数totalTiles总 tile 数 A1 × tileA0OutertiletilesPerCore每核 tile 数tileusedCoreNums实际使用核数核tileA0OuterA0 方向 tile 总数tiletileA0Len / tileA0Tail主块/尾块 A0 长度元素数binAddRFactorR 方向单 bin 长度默认 128元素数binAddRLoopR 方向完整 bin 数binbinAddRTotalLoopR 方向 bin 总数binbinAddRTailR 方向尾 bin 长度元素数binAddBasicBlockLoop二分树基础块数块binAddMainFoldCount需折叠块数块binAddCacheBufferCountcache buffer 数量个binAddResultCacheID最终结果 cache 索引索引6. Host tiling 参数计算方法使用 softmax_v2_tiling.h 中的TilingAraRecomputesoftmax_tiling::CaseShape shape{a1, r, a0, dtypeCode}; softmax_tiling::PlatformParam plat{ubSize, numBlocks}; SoftmaxV2ARARecomputeTilingData td; int64_t blockDim softmax_tiling::TilingAraRecompute(shape, plat, td);关键计算binAddRFactor 128默认binAddRTotalLoop CeilDiv(R, binAddRFactor)binAddBasicBlockLoop FindNearestPower2(binAddRTotalLoop)binAddCacheBufferCount 64 - clz(binAddBasicBlockLoop)二分树层数binAddResultCacheID GetCacheID(binAddBasicBlockLoop - 1)7. CopyIn、归约、归一化和 CopyOut 流程三阶段处理每 tile 独立执行Process逐 tile: Step 1 — CalcReduceMax: for 每个 R bin: CopyInX → 逐 A0 列 VF 求 max与 xMaxLocal 合并 Step 2 — CalcReduceSum: for basicBlockIdx in [0, binAddBasicBlockLoop): a. ProcessMainBlock: CopyIn 主 bin → sub max → exp → 写入 yMain b. 若需折叠: ProcessFoldBlock: CopyIn 折叠 bin → sub max → exp → Add 到 yMain c. ProcessSummation: NlastReduceSum(yMain) → UpdateCache binAddBasicBlockLoop 0 时: 单 bin 直接处理 最终: 从 cache[resultCacheID] 拷出 xSum Step 3 — 输出: for 每个 R bin: CopyInX → CalcOutputsub max, exp, div sum, Cast→ CopyOutY8. NlastReduceSum、bin fold 和 cache 参数NlastReduceSumNlastReduceSum是跨 R行方向对 A0列做规约求和的核心函数实现在 softmax_v2_base.h 中R ≤ 8NlastReduceSumSmallR使用NlastDichotomyAddRSize模板做编译期二分展开。NlastDichotomyAdd2基例2 行Add。递归展开NlastDichotomyAddN调用NlastDichotomyAdd(N1)/2和NlastDichotomyAddN/2。TailCount模板参数处理尾块不足 8 行时的精确掩码。R 8NlastReduceSumLargeRTailCount按 8 行分组COMPRESSION8做折叠。FindNearestPower2(rSize)确定 foldPoint。mainFold8 行折叠、tailFold尾块折叠、unFold展开三阶段。递归调用NlastReduceSumSmallR做最终归约。bin foldMainBlock / FoldBlock 配对ProcessMainBlock载入主 binSub(maxReg)→Exp→ 写入 yMain。ProcessFoldBlock载入折叠 binSub(maxReg)→Exp→Add到 yMain。配对后一次NlastReduceSum归约次数减半。cache 参数binAddCacheBufferCount二分树层数 1决定 cacheBuffer 大小。binAddResultCacheIDGetCacheID(binAddBasicBlockLoop - 1)指向最终合并结果。UpdateCache按cacheID层级将当前 sum 累加到 cache 对应位置。9. 主块、尾块、对齐和数值精度处理主块/尾块tileA0Len为主块 A0 长度tileA0Tail为尾块。binAddRTail为 R 方向尾 bin 长度。折叠尾 bin当basicBlockIdx binAddMainFoldCount且binAddRTail 0且binAddRTail ! binAddRFactor时折叠块使用binAddRTail长度。对齐A0 按a0TileBaseFP328/FP1616对齐dstStride按BLOCK_SIZE32B对齐。数值精度FP16/BF16 通过LoadTensorForDtypeT1升至 FP32max/sum 用 FP32 维护输出Cast降回。10. 独立 Kernel 入口示例#include softmax_v2_tiling_data.h #include dav310/softmax_v2_base.h #include dav310/softmax_v2_ara_recompute.h using namespace SoftmaxV2Ops; extern C __global__ __aicore__ __vector__ void softmax_ara_recompute_fp32(GM_ADDR x, GM_ADDR y, SoftmaxV2ARARecomputeTilingData tiling) { TPipe pipe; SoftmaxV2ARARecomputefloat, float op(tiling); op.Init(x, y, pipe); op.Process(); }11. 不适用场景和可选替代模板不适用场景推荐替代R×tileA0 可载入 UBARA FullLoad无重读A0 12DAR Recompute需要在线更新避免 3 次读入ARA Online2 次读入带宽极度受限评估 FullLoad 是否可用3× 重读代价高【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考