CANN/cannbot-skills: SoftmaxV2 AR重计算模板

SoftmaxV2 AR Recompute 模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 模板用途针对 2D[A1, R]布局、R 超出 UB 容量的 Softmax 场景。将 R 切分为多个 chunk通过三阶段重读 GM求 max → 求 sum → 输出完成 Softmax使用 UB 间二分累加树UpdateCache高效合并跨 chunk 的局部 sum。2. 输入输出布局和 Softmax 轴输入x[A1, R]行优先Softmax 沿 R 轴计算。输出y[A1, R]与输入同 shape、同 dtype。3. 适用 shape、dtype 和 UB 条件条件要求输入维数2DA0 1R 大小R 超 UB 单次载入容量dtypeFP32、FP16、BF16UB 约束预留 max(32B) sum(32B) cache(2048B) 后三缓冲 xQueue 双缓冲 yQueue tmpBuffer4. 类名、构造函数和 Init 接口namespace SoftmaxV2Ops; template typename Tx, typename Ty class SoftmaxV2ArRecompute : public SoftmaxV2OpsBase { public: __aicore__ inline SoftmaxV2ArRecompute(TPipe* pipe); __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const SoftmaxV2ArRecomputeTilingData* tilingData); __aicore__ inline void Process(); };构造函数接收TPipe*。Init计算每核行数currentRowBlock_、resultCacheID_初始化 xQueue三缓冲、yQueue双缓冲及 xTmp/xMax/xSum/cache buffer。Process逐行执行三阶段计算。5. 对应的 TilingData 和字段说明对应SoftmaxV2ArRecomputeTilingData定义在 softmax_v2_tiling_data.h。字段含义单位aA1 维总长度行数元素数rR 维实际长度元素数ubFactorR 方向单 chunk 长度元素数ubFactorTailR 方向尾块长度元素数aBlockFactor每核处理行数行aLoopCountCeilR 方向 chunk 总数chunkbasicBlockLoop二分树基础块数≤ aLoopCountCeil 的最大 2^k块mainFoldCount需折叠块数 FloorDiv(R,ubFactor) - basicBlockLoop块6. Host tiling 参数计算方法使用 softmax_v2_tiling.h 中的TilingArRecomputesoftmax_tiling::CaseShape shape{a1, r, 1, dtypeCode}; softmax_tiling::PlatformParam plat{ubSize, numBlocks}; SoftmaxV2ArRecomputeTilingData td; int64_t blockDim softmax_tiling::TilingArRecompute(shape, plat, td);关键计算ubFactor min(R, (ubFlexible / baseFactor) 按 lcm 对齐)其中 baseFactor dtypeSize×3 dtypeSize×2 4aLoopCountCeil CeilDiv(R, ubFactor)basicBlockLoop FindNearestPower2(aLoopCountCeil)mainFoldCount FloorDiv(R, ubFactor) - basicBlockLoop7. CopyIn、归约、归一化和 CopyOut 流程三阶段处理每行独立执行Process逐行 rowIdx: Step 1 — 求 max: for 每个 R chunk: CopyIn chunk → CalculateMaxVFVF 求 max与 running max 合并 Step 2 — 求 sum二分累加树: for basicBlockIdx in [0, basicBlockLoop): a. CopyIn 主块 → MainBlockCastSubExpVFsub max, exp, 写入 xTmp b. 若需折叠: CopyIn 折叠块 → FoldBlockCastSubExpVFsub max, exp, Add 到 xTmp c. ReduceSum(xTmp) → UpdateCache按 cacheID 层级合并 basicBlockLoop 0 时: 单 chunk 直接 ReduceSum Step 3 — 输出: for 每个 R chunk: CopyIn chunk → CalculateOutVFsub max, exp, div sum, Cast 回原精度→ CopyOut GM8. 重计算与二分累加树优化三阶段重读输入读 3 次max/sum/output换取 UB 空间——R 可远超 UB 容量。MainBlock / FoldBlock 配对主块和折叠块的 exp 结果在 xTmp 上Add合并后一次ReduceSum归约次数减半。UpdateCache 二分树跨 chunk 的局部 sum 按GetCacheID(basicBlockIdx)层级二分合并O(log N) 层而非 O(N) 串行。GetCacheID(idx)popcount(idx ^ (idx1)) - 1确定当前块在二分树中的位置。resultCacheIDGetCacheID(basicBlockLoop - 1)指向最终合并结果。三缓冲 xQueue支持 Main/Fold 双 chunk 并行载入。9. 主块、尾块、对齐和数值精度处理主块/尾块ubFactor为主 chunk 长度ubFactorTail为尾 chunkR 不能整除时。尾块在 Step 1/3 中通过条件判断使用ubFactorTail。折叠尾块当basicBlockIdx mainFoldCount且ubFactorTail 0时折叠块使用ubFactorTail长度。对齐ubFactor按Lcm(blockSize, Lcm(blockSize/dtypeSize, blockSize/dtypeSize))对齐。数值精度FP16/BF16 通过LoadTensorForDtypeTIn升至 FP32 计算max/sum 用 FP32 维护输出通过Cast降回。xToFp32_/yToFp32_编译期常量控制转换路径。10. 独立 Kernel 入口示例#include softmax_v2_tiling_data.h #include dav310/softmax_v2_base.h #include dav310/softmax_v2_ar_recompute.h using namespace SoftmaxV2Ops; extern C __global__ __aicore__ __vector__ void softmax_ar_recompute_fp32(GM_ADDR x, GM_ADDR y, SoftmaxV2ArRecomputeTilingData tiling) { TPipe pipe; SoftmaxV2ArRecomputefloat, float op(pipe); op.Init(x, y, tiling); op.Process(); }11. 不适用场景和可选替代模板不适用场景推荐替代R 可载入 UBAR FullLoad无重读带宽更优R 极小≤ 16/32AR SmallR转置向量化3D [A1,R,A0] 布局A01ARA Recompute带宽极度受限、R 适中评估 FullLoad 是否可用3× 重读代价高【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考