SoftmaxV2 AR SmallR 模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 模板用途针对 2D[A1, R]布局、归约轴 R 极小FP32 ≤ 16 / FP16/BF16 ≤ 32的 Softmax 场景。通过将数据转置为(R, A1)布局使短归约轴转到外层、长伴生轴转为向量化方向将 VReg 利用率从 R/64 提升到接近 100%。2. 输入输出布局和 Softmax 轴输入x[A1, R]行优先Softmax 沿 R 轴最后一轴计算。输出y[A1, R]与输入同 shape、同 dtype。内部转置为(R, A1)布局进行 VF 计算输出前转置回(A1, R)。3. 适用 shape、dtype 和 UB 条件条件要求输入维数2DA0 1R 大小FP32 ≤ 16FP16/BF16 ≤ 32R ≤ 2×VL_FP32 128 时可转置向量化dtypeFP32、FP16、BF16UB 约束tileA0Len * rAligned * (dtypeSize*2 4) * 2 tileA0Len * 4≤ UB4. 类名、构造函数和 Init 接口namespace SoftmaxV2Ops; template typename Tx, typename Ty class SoftmaxV2ArSmallR { public: __aicore__ inline SoftmaxV2ArSmallR(TPipe* pipe); __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const SoftmaxV2ArSmallRTilingData* tilingDataIn); __aicore__ inline void Process(); };构造函数接收TPipe*保存 pipe 指针。Init接收输入 x、输出 y 的 GM 地址和 TilingData 指针初始化 GlobalTensor 和 UB Buffer。Process执行完整 Softmax 计算。5. 对应的 TilingData 和字段说明对应SoftmaxV2ArSmallRTilingData定义在 softmax_v2_tiling_data.h。字段含义单位totalA0Len转置后 A0 维长度 原 A1元素数totalRLen归约轴 R 实际长度元素数totalTilesA0 方向总 tile 数tiletilesPerCore每核 tile 数tiletileA0Len主块 A0 tile 长度元素数tileA0Tail尾块 A0 长度元素数rTileBaseR 转置基础块FP328/FP1616元素数rAlignedR 对齐后长度元素数6. Host tiling 参数计算方法使用 softmax_v2_tiling.h 中的TilingArSmallRsoftmax_tiling::CaseShape shape{a1, r, 1, dtypeCode}; softmax_tiling::PlatformParam plat{ubSize, numBlocks}; SoftmaxV2ArSmallRTilingData td; int64_t blockDim softmax_tiling::TilingArSmallR(shape, plat, td);关键计算rTileBase FP32 ? 8 : 16rAligned CeilDiv(R, rTileBase) × rTileBasetileA0Len a0TileNum × VL_FP32受 UB 容量约束totalTiles CeilDiv(A1, tileA0Len)7. CopyIn、归约、归一化和 CopyOut 流程Process: 1. preload 第一个 tileCopyInAndTransPoseMultiCopy 转置 (A1,R)→(R,A1) 2. for 每个 tile: a. CalcMaxSubExp: 逐 A0 列求 maxsub maxexp写入 tmpLocal 和 tmpLocal2 b. CopyInAndTransPose: 预载下一个 tile双缓冲 c. CalcReduceSum: 沿 R 轴 ReduceSumPattern::Reduce::RA d. CalcOutput: 逐列 div sumFP16/BF16 做 Cast e. CalcTranspose: TransDataTo5HD 转置回 (A1,R) f. CopyOutY: DataCopyPad 写回 GM8. 转置优化MultiCopy 转置使用MultiCopyParamsTx, 2MultiCopyLoopInfo2在 CopyIn 阶段直接完成(A1,R)→(R,A1)转置避免额外转置 pass。TransDataTo5HD输出前使用TransDataTo5HD将(R,A1)转回(A1,R)FP32 和 FP16/BF16 分别使用不同的 block 参数。VF 向量化转置后长轴 A1 沿 VReg 方向每个 VReg 可处理 64 个 FP32 元素利用率接近 100%。9. 主块、尾块、对齐和数值精度处理主块/尾块tileA0Len为主块长度tileA0Tail为最后一个 tile 的实际长度。尾块通过UpdateMask精确掩码处理。对齐R 按rTileBaseFP328/FP1616向上对齐A0 按VL_FP3264对齐。数值精度FP16/BF16 输入在计算前Cast升至 FP32计算完成后Cast降回原精度。使用castTraitFp16ToFp32/castTraitFp32ToFp16控制饱和模式和舍入。10. 独立 Kernel 入口示例#include softmax_v2_tiling_data.h #include dav310/softmax_v2_base.h #include dav310/softmax_v2_ar_small_r.h using namespace SoftmaxV2Ops; extern C __global__ __aicore__ __vector__ void softmax_ar_small_r_fp32(GM_ADDR x, GM_ADDR y, SoftmaxV2ArSmallRTilingData tiling) { TPipe pipe; SoftmaxV2ArSmallRfloat, float op(pipe); op.Init(x, y, tiling); op.Process(); }11. 不适用场景和可选替代模板不适用场景推荐替代R 2×VL_FP32128AR FullLoadR 可载入 UB或 AR RecomputeR 超 UB3D [A1,R,A0] 布局A01ARA FullLoad 或 ARA Recompute【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考