SoftmaxV2 ARA Online Softmax 模板【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 模板用途针对 3D[A1, R, A0]布局的独立 Softmax 场景使用在线online算法沿 R 分 chunk 逐块搬入为每个 A0 列维护 running max 和 running sum。相比 ARA Recompute输入读 3 次max/sum/outputOnline Softmax 将 max 与 sum 融合为单遍在线更新输入只需读 2 次1 遍在线 maxsum1 遍输出减少一次 R 全量搬入。重要本模板是独立的 Online Softmax不包含QK^T、P×V、O_acc、causal mask 或 FlashAttention 的 Cube/Vector 融合。它与 融合 Attention Online Softmax 设计 只共享在线 max/sum 数学基础不是同一个实现。2. 输入输出布局和 Softmax 轴输入x[A1, R, A0]行优先Softmax 沿 R 轴计算。输出y[A1, R, A0]与输入同 shape、同 dtype。输入为已经生成的 Softmax 数据输出为完整 Softmax 概率。3. 适用 shape、dtype 和 UB 条件条件要求输入维数3DA0 1R 大小任意R 分 chunk 处理不受 UB 限制dtypeFP32、FP16、BF16UB 约束双缓冲 xQueue/yQueue xMaxBuf xSumBuf按 rChunkFactor 分 chunk4. 类名、构造函数和 Init 接口namespace SoftmaxV2Ops; template typename T1, typename T2 class SoftmaxV2ARAOnline : public SoftmaxV2OpsBase { public: __aicore__ inline SoftmaxV2ARAOnline(); __aicore__ inline SoftmaxV2ARAOnline(const SoftmaxV2ARAOnlineTilingData* tilingDataIn); __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, TPipe* pipeIn); __aicore__ inline void Process(); };构造函数无参默认或接收 TilingData 指针。Init初始化 GM/UB bufferxQueue/yQueue双缓冲大小 tileA0Len × rChunkFactor xMaxBuf/xSumBuf。Process按 tile 遍历每 tile 执行在线 maxsum 和输出两遍。5. 对应的 TilingData 和字段说明对应SoftmaxV2ARAOnlineTilingData定义在 softmax_v2_tiling_data.h。字段含义单位totalRLen / totalA0LenR/A0 维总长度元素数totalTiles总 tile 数 A1 × tileA0OutertiletilesPerCore每核 tile 数tileusedCoreNums实际使用核数核tileA0OuterA0 方向 tile 总数tiletileA0Len / tileA0Tail主块/尾块 A0 长度元素数rChunkFactorR 方向单 chunk 长度默认 128元素数rChunkTotalLoopR 方向 chunk 总数chunkrChunkTailR 方向尾 chunk 长度元素数6. Host tiling 参数计算方法使用 softmax_v2_tiling.h 中的TilingAraOnlinesoftmax_tiling::CaseShape shape{a1, r, a0, dtypeCode}; softmax_tiling::PlatformParam plat{ubSize, numBlocks}; SoftmaxV2ARAOnlineTilingData td; int64_t blockDim softmax_tiling::TilingAraOnline(shape, plat, td);关键计算rChunkFactor min(128, R)rChunkTotalLoop CeilDiv(R, rChunkFactor)rChunkTail R - (rChunkTotalLoop - 1) × rChunkFactorUB 占用每个 A0 列2 × rChunkFactor × dtypeSize 2 × rChunkFactor × 4 4 4字节7. 在线更新公式沿 R 分 chunk按 A0 列独立维护 running max (m) 和 running sum (l)m_new max(m_old, chunk_max) corr exp(m_old - m_new) l_new l_old * corr sum(exp(x - m_new))最终 Softmax 输出softmax(x) exp(x - m_final) / l_final与 ARA Recompute 三次读取输入的区别特性ARA RecomputeARA Online输入读取次数3 次max/sum/output2 次在线 maxsum / outputmax 与 sum 计算分离先 max 全量再 sum 全量融合单遍在线更新cache buffer需要二分累加树不需要数值稳定性Safe Softmax先减 maxSafe Softmax在线减 running max带宽3× R 全量读2× R 全量读8. CopyIn、归约、归一化和 CopyOut 流程Process逐 tile: 第一遍 — CalcOnlineMaxSum: 初始化 running max -inf, running sum 0 for 每个 R chunk: CopyInX → 逐 A0 列: 1. 求 chunk 内局部 max (mLocal) 2. m_new max(m_old, mLocal) 3. corr exp(m_old - m_new) 4. chunk_sum Σ exp(x - m_new) 5. l_new l_old * corr chunk_sum 6. 更新 xMaxLocal m_new, xSumLocal l_new 第二遍 — 输出: for 每个 R chunk: CopyInX → CalcOutput: sub max → exp → div sum → Cast 回原精度 CopyOutY → GM9. 主块、尾块、FP32 running 状态和数值稳定性主块/尾块tileA0Len为主块 A0 长度tileA0Tail为尾块。rChunkTail为 R 方向尾 chunk 长度。对齐A0 按a0TileBaseFP328/FP1616对齐。FP32 running 状态running max 和 running sum 始终用 FP32 维护xMaxBuf/xSumBuf为float类型即使输入/输出为 FP16/BF16确保数值稳定性。数值稳定性每步减去 running maxm_new避免 exp 溢出。corr exp(m_old - m_new)≤ 1因 m_new ≥ m_old修正历史 sum。第一遍在线计算最终 max 和 sum第二遍用最终值做归一化。10. 独立 Kernel 入口示例#include softmax_v2_tiling_data.h #include dav310/softmax_v2_base.h #include dav310/softmax_v2_ara_online.h using namespace SoftmaxV2Ops; extern C __global__ __aicore__ __vector__ void softmax_ara_online_fp32(GM_ADDR x, GM_ADDR y, SoftmaxV2ARAOnlineTilingData tiling) { TPipe pipe; SoftmaxV2ARAOnlinefloat, float op(tiling); op.Init(x, y, pipe); op.Process(); }Host tiling 示例#include softmax_v2_tiling.h softmax_tiling::PlatformParam plat; plat.ubSize ubSize; // 运行时查询 plat.numBlocks numBlocks; // 运行时查询 softmax_tiling::CaseShape shape; shape.a1 4; shape.r 2049; shape.a0 256; shape.dtypeCode 0; // FP32 SoftmaxV2ARAOnlineTilingData td; int64_t blockDim softmax_tiling::TilingAraOnline(shape, plat, td); // blockDim 为 启动核数11. 本模板不包含的内容本模板是独立的 Online Softmax不包含QK^Tscore 计算P×V概率与 Value 的乘积O_acc输出累加器causal mask因果掩码Attention workspaceFlashAttention 的 Cube/Vector 融合如需了解融合 Attention 中的 Online Softmax 设计参见 融合 Attention Online Softmax 设计。两者只共享在线 max/sum 数学基础不是同一个实现。12. 不适用场景和可选替代模板不适用场景推荐替代A0 12DAR RecomputeR×tileA0 可载入 UB 且带宽充裕ARA FullLoad单次载入R 超 UB 且不需要在线更新ARA Recompute三阶段重读【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考