Softmax — Tiling 实践案例【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills以 Softmax 为例展示如何将算子数学流程映射到 通用五模板。 Agent 推导 LayerNorm、RMSNorm 等算子时可复用相同结构。1. 算子概述Softmax 沿 R 轴计算output[i] exp(input[i] - max(R)) / sum(exp(input[R] - max(R)))等价于四步流水步骤操作沿 R 轴类型①ReduceMax → max归约②Sub Exp → exp_shifted逐元素③ReduceSum → sum归约④Div → output逐元素关键特征步骤 ②③ 需要步骤 ① 的结果max步骤 ④ 需要步骤 ③ 的结果sum。 Recompute 模式下 R 无法一次载入需分轮重读原数据。2. 形状准备Softmax 的归约轴为 R合轴后得到(A1, R, A0)原始 shapeaxes合轴结果[batch, seq, dim]axis1axis1A1batch, Rseq, A0dim[batch, seq]axis1axis1A1batch, Rseq, A01 → AR 族[N]axis0axis0A11, RN, A01 → AR 族3. 模板选型Softmax 视角Softmax 使用通用决策树无额外分支。以下说明各模板下 Softmax 的具体行为。AR 族A01如[batch, seq]A0 1 │ ├─ R 极小≤ 约 16~32 │ → AR-SmallR │ 转置为 (R, batch)沿 batch 切 tile │ 每个 tile 内完成完整 Softmax 四步 │ ├─ R 可整段载入 UB │ → AR-FullLoad │ R 行一次载入batch 方向多核切分 │ max / exp / sum / div 均在 UB 内完成数据只读一次 │ └─ R 超出 UB → AR-Recompute R 方向分块每块做 max exp partial sum 块间合并 max 和 sum最后分轮重读原数据做 divARA 族A01如[batch, seq, dim]A0 1 │ ├─ R 可整段载入 UB │ → ARA-FullLoad │ 沿 dim(A0) 切 tile每个 tile 内 R 行全载 │ R ≤ 8: 直接多行累加求 max/sum │ R 8: 用二分累加做行归约 │ └─ R 超出 UB → ARA-Recompute 沿 dim(A0) 切 tileR 按 128 行分 bin 每个 bin 内做 partial max/sum跨 bin 二分合并4. UB 预算修正相对标准归约Softmax 在通用公式基础上UB 内同时存在多份 bufferAR-FullLoad每行 UB 开销 r_align × (输入×2 输出×2 FP32中间量×2) 固定预留 1024 512 rows_per_ub (ub_size - 1536) / (r_align × (4 4×dtype_bytes)) rows_per_core ceil(A1 / core_num) rows_per_ub min(rows_per_ub, rows_per_core) FullLoad 条件: rows_per_ub ≥ 1对比标准 ReduceSum分母多了 FP32 中间量和双倍 buffer输入/输出各 2 份。AR-Recompute可用 UB ub_size - 2112 // max_buf(32) sum_buf(32) binary_cache(2048) 每元素开销 输入×3 输出×2 FP32×1 // 3 份输入: 原数据 exp 结果 重读缓冲 // Recompute 需分轮重读原数据计算 exp r_chunk floor(可用UB / 每元素开销)32B 对齐AR-SmallRr_tile_unit 8(FP32) / 16(FP16) r_align ceil(R, r_tile_unit) × r_tile_unit per_tile r_align × (输入×2 输出×2 FP32×2) max_a1_tiles ub_size / (64 × (per_tile 4))ARA-FullLoad / ARA-RecomputeARA-FullLoad 每 tile: per_tile R × (输入×2 FP32×2 FP32) 8 ARA-Recompute 每 tile: r_bin_size 128 per_tile 128 × (输入×2 FP32) FP32 × (11 cache_layers)5. 典型案例案例 A[32, 512]FP32axis1合轴: A132, R512, A01 → AR 族R512 超出 SmallR 阈值检查 FullLoad 预算假设 ub_size248KBrows_per_ub ≥ 1 →AR-FullLoad多核: rows_per_core ceil(32/64) 1used_core_num 32案例 B[8, 65536]FP32axis1合轴: A18, R65536, A01 → AR 族R 远超 UB 容量 →AR-Recomputer_chunk 由可用 UB 和每元素开销决定块间用二分累加合并 partial max 和 partial sum案例 C[2, 64, 128]FP16axis1合轴: A12, R64, A0128 → ARA 族R64 可全载 →ARA-FullLoadR 8启用 BinaryAdd 做行归约a0_tile_len 由 UB 容量和多核均衡共同决定6. 推导 Norm 类算子的提示LayerNorm / RMSNorm 与 Softmax 结构类似沿 R 轴做 mean → var → normalize对比项SoftmaxLayerNorm归约次数max summean var两个关联统计量中间 bufferexp 缓冲去均值缓冲、方差缓冲Recompute 增强重读原数据建议启用Welford Online模板选择通用五模板相同决策树修正 UB 预算Agent 步骤列出 Norm 算子沿 R 轴的计算步骤数清 UB 内同时需要的 buffer 份数套用 通用 UB 预算公式若有两个关联统计量且走 Recompute启用 Welford7. 选型自检输入已合轴为 (A1, R, A0)dtype 已确定FP32/FP16/BF16平台参数 ub_size / core_num 已获取AR 族: 按 R 从小到大尝试 SmallR → FullLoad → RecomputeARA 族: 按 UB 预算选 FullLoad → RecomputeSoftmax Recompute 模式确认需重读原数据尾块: 检查 a1_tile_tail / r_chunk_tail / a0_tile_tailused_core_num ≥ 18. 关联文档通用五模板决策树TilingData 字段语义参考实现简化预算【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考