Reduction 归约类算子高性能实现索引【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skillsReduce 族算子Softmax、LayerNorm、RMSNorm、ReduceSum、ReduceMax 等在 Ascend Vector 核心上的高性能实现沉淀为一组可复用 Kernel 模板。本指南是这些模板的入口介绍每个模板解决的问题、关键技术与适用条件以及如何适配到同族其他算子如 Norm。本指南是优秀实践目录而非选型决策树每个模板列出其适用条件与解决的问题由实现者根据实际 shape、UB 容量与带宽情况判断采用。Reduce 族算子形态多样不存在普适的唯一最优理解各模板的权衡是正确适配的前提。文档结构guide.md本文件 └── templates/ ├── usage_guide.md 模板集成指南 ├── softmax_v2_tiling_data.template 六套 TilingData 公共定义→ .h ├── softmax_v2_tiling.template 六个独立 Host tiling 参考实现→ .h ├── dav310/ │ ├── kernel_utils.template Device 侧公共工具模板→ .h │ ├── softmax_v2_base.template 公共基类 SoftmaxV2OpsBase→ .h │ ├── softmax_v2_ar_small_r.template AR-SmallR Kernel→ .h │ ├── softmax_v2_ar_full_load.template AR-FullLoad Kernel→ .h │ ├── softmax_v2_ar_recompute.template AR-Recompute Kernel→ .h │ ├── softmax_v2_ara_full_load.template ARA-FullLoad Kernel→ .h │ ├── softmax_v2_ara_recompute.template ARA-Recompute Kernel→ .h │ └── softmax_v2_ara_online.template ARA-Online Kernel→ .h ├── softmax_v2_ar_small_r.md AR-SmallR 说明文档 ├── softmax_v2_ar_full_load.md AR-FullLoad 说明文档 ├── softmax_v2_ar_recompute.md AR-Recompute 说明文档 ├── softmax_v2_ara_full_load.md ARA-FullLoad 说明文档 ├── softmax_v2_ara_recompute.md ARA-Recompute 说明文档 ├── softmax_v2_ara_online.md ARA-Online 说明文档 ├── fused_attention_online_softmax_design.md 融合 Attention Online Softmax 设计 └── state_resident_design.md 融合算子 UB 状态常驻设计阅读顺序先读本指南「通用范式」「公共框架」理解骨架 → 读对应模板.md了解流程与关键技术 → 落代码时对照同名.template使用时去掉.template后缀转成.h见 usage_guide.md 第 0 节→ 集成时读 usage_guide.md → 在线场景读fused_attention_online_softmax_design.md/state_resident_design.md。⚠️集成前必读usage_guide.md 说明了如何将模板集成到实际项目中包括构建配置、kernel 入口写法、host 调用方式、常见编译问题及解决方案。Reduce 计算的通用范式Reduce 族算子均可拆为两段模板的差异主要在如何切分 R 以适配 UB而两段数学是可替换的① 归约统计reduce statistic沿 R 轴求一个标量/向量统计量 Softmax → max、sum(exp) LayerNorm → mean、variance RMSNorm → mean(square) ReduceSum → sum ReduceMax → max ② 归一化/输出normalize用统计量逐元素变换 Softmax → (x - max) → exp → div sum LayerNorm → (x - mean) / sqrt(var eps) → scale shift RMSNorm → x / sqrt(mean_sq eps) → scale正因如此6 个 Softmax 模板的骨架切分 流水 归约统计的组织方式对 Norm 等同族算子通用适配时只需替换两段数学见下文「适配同类算子」。公共框架与工具所有模板继承 dav310/softmax_v2_base.template转换后softmax_v2_base.h中的SoftmaxV2OpsBase提供跨模板复用的 VF 工具工具说明CastTraitMicroAPI Cast 必需配置饱和模式、舍入、掩码合并castTraitFp16ToFp32/castTraitFp32ToFp16CastToFp32FromT/CastFromFp32ToT载入即升 FP32、输出即降回原精度FP32 直通、FP16/BF16 自动 CastUpdateCache跨 chunk/bin 二分累加树核心按cacheID层级归并局部 sumrecompute 类模板共用GetCacheID/FindNearestPower2二分树层级计算确定配对关系与根节点位置NlastReduceSumN-last 方向规约求和小 R 用NlastDichotomyAdd编译期二分展开大 R 用NlastReduceSumLargeR8 行分组折叠CopyIn/CopyOut带 stride 的批量 DMArows×cols支持非对齐尾块常量VL_FP3264256B VReg / 4B float、CONST_EIGHT8、BLOCK_SIZE32UpdateCache和NlastReduceSum是 recompute 系模板AR-Recompute / ARA-Recompute的关键把逐 chunk 串行累加 sum变为二分树对数层合并显著降低长 R 的归约串行延迟。离线计算模板离线模板针对独立 Kernel的 Reduce 算子输入已完整存在于 GM按 R 能否载入 UB 与输入维数分为 5 个。各模板互相独立Host 侧 Tiling 根据 shape 与 UB 容量选用。注下表适用条件是模板的设计前提非强制决策规则。边界 shape如 R 刚好接近 UB 上限可由实现者权衡带宽与空间。模板输入R 与 UB 关系核心策略代码说明AR-SmallR2D[A1,R]R 极小FP32≤16/FP16≤32转置为(R,A1)沿 A1 向量化.template.mdAR-FullLoad2D[A1,R]R 全量载入 UB多行批量载入单 pass VF 计算.template.mdAR-Recompute2D[A1,R]R 超出 UBR 切 chunk3 阶段重读 GM 二分累加树.template.mdARA-FullLoad3D[A1,R,A0]R 全量载入 UB沿 A0 切 tile跨 R 逐列 VF.template.mdARA-Recompute3D[A1,R,A0]R 超出 UBR 切 bin3 阶段跨 bin 二分折叠.template.md各模板的「解决的问题 / 适用条件 / 执行流程 / 关键技术 / 关键代码索引」详见对应.md文档。维度命名约定A1最外层非归约轴行方向2D 中即行数。R归约轴reduce axisSoftmax/Norm 沿此求统计量。A03D 中 R 内侧的非归约轴列方向与 R 共同决定一个 tile 的数据量。AR2D(A1, R)ARA3D(A1, R, A0)。在线计算模板独立 ARA Online Softmax模板输入核心策略代码说明ARA-Online3D[A1,R,A0]沿 R 分 chunk 在线更新 running max/sum输入读 2 次.template.mdARA Online Softmax 是独立的 Online Softmax不包含QK^T、P×V或 FlashAttention 融合。与 ARA Recompute输入读 3 次相比减少一次 R 全量搬入。融合 Attention Online Softmax 设计文档解决的问题关键思想fused_attention_online_softmax_design.mdFlashAttention 中 Softmax 需 O(S²) 中间内存逐 tile 生成QK^T维护 running max/sum融合P×V内存降至 O(S)state_resident_design.md融合算子跨 S2 循环重复分配状态 buffer状态 buffer 一次性常驻 UB 双缓冲索引重要区分softmax_v2_ara_online.md 是独立的 Online Softmax Kernel 实现输入读 2 次输出完整概率。fused_attention_online_softmax_design.md 是 FlashAttention 融合场景的设计参考逐 tile 生成 score 即消费融合P×V通常不生成完整概率矩阵。两者共享在线 max/sum 数学基础但不是同一个实现。当前目录未提供与融合设计完全对应的 FlashAttention Kernel。关键技术总览以下技术横跨多个模板是 Reduce 族高性能实现的公共经验技术涉及模板核心要点DOUBLE/TRIPLE BUFFER 乒乓全部TQue多缓冲让 CopyIn/Compute/CopyOut 重叠recompute 用 3 缓冲支持 Main/Fold 双 chunk 并行载入MicroAPI VF__VEC_SCOPE__全部RegTensorfloatMaskReg直接操作 256B VReg逐VL_FP3264块处理尾块用UpdateMask精确掩码FP16/BF16 ↔ FP32 计算全部载入即升 FP32 保证数值稳定输出即降回CastTrait控饱和/舍入批量 strided DMA全部DataCopyPadDataCopyExtParamsblockCount一次搬多行srcStride/dstStride跨维跳跃对齐到 32B block二分累加树UpdateCacheAR/ARA-Recompute跨 chunk/bin 的局部 sum 按GetCacheID层级二分合并O(log N) 层而非 O(N) 串行NlastReduceSumARA-FullLoad/RecomputeN-last 方向规约小 R 用NlastDichotomyAdd编译期二分展开大 R 用 8 行分组折叠MainBlock/FoldBlock 配对AR/ARA-Recompute两 chunk 的 exp 在 tmp 上Add合并后一次 ReduceSum归约次数减半布局转置向量化AR-SmallR短归约轴转置到外层长伴生轴转成向量化方向VReg 利用率从 R/64 拉到 ~100%BinaryAddVFARA-FullLoadR8 时按 8 行分组二分折叠TwoRowAddWithTail处理余数尾块按 R 大小分支归约ARA-FullLoadR≤2/≤4/≤8 用 VFAdd直累R8 用BinaryAddVF二分折叠在线 max/sum 更新ARA-Online沿 R 分 chunkrunning max/sum 单遍在线更新减少一次输入读取适配同类算子Softmax 模板的骨架切分 流水 归约统计组织对同族算子通用适配只需替换两段数学。下面以LayerNorm / RMSNorm为例说明映射关系。适配要点维度映射Norm 的归约轴hidden_size / norm_size对应模板的 R其余轴对应 A12D或 A1A03D。2D LayerNorm[N, H]沿 H 归约 → AR 模板RH, A1N。3D LayerNorm[B, S, H]沿 H 归约 → 可视为 ARARH, A0S, A1B或展平为 AR。数学替换把 Softmax 的max/sub/exp/sum/div替换为 Norm 的统计量计算见下表。阶段数调整Softmax 是三段max → exp-sum → div。LayerNorm 可两遍扫描mean → var或 Welford 单遍recompute 模板的阶段①可同时累加 mean阶段②算 var。RMSNorm 单段即可平方累加 → div sqrt比 Softmax 更简单。affine 参数Norm 通常带 gamma/betascale/shift可在输出阶段追加MulsAdd。精度mean/variance 建议用 FP32 累加与 Softmax 的 FP32 中间一致即使输入输出为 FP16。数学替换对照算子阶段①reduce statistic阶段②sum/second pass阶段③normalize/outputSoftmaxglobal maxexp 累加求 sumsub max→exp→div sumLayerNorm两遍累加求 meansub mean→平方→累加求 varsub mean→div sqrt(vareps)→scaleshiftRMSNorm平方累加可并入①div sqrt(mean_sqeps)→scaleReduceSum累加求 sum—直接输出ReduceMaxmax 累加—直接输出模板选择参考Norm 变体Rnorm_size特征倾向模板LayerNorm/RMSNormhidden 小极小AR-SmallR转置向量化收益高LayerNorm/RMSNormhidden 中等全量载入AR-FullLoad / ARA-FullLoadLayerNorm/RMSNormhidden 很大超出 UBAR-Recompute / ARA-Recompute以上是特征—模板对应关系非决策树。实际选择还需结合 A1/A0 规模、UB 预算、是否需 affine、是否与 matmul 融合等综合判断。Softmax 独有优化在线融合以下两个设计文档主要服务于 Softmax 在 Attention 中的融合场景对独立 Norm 算子参考价值有限列出供完整性文档说明fused_attention_online_softmax_design.md融合 Attention Online Softmax 设计逐 tile 生成QK^Trunning max/sum P×V融合内存 O(S²)→O(S)state_resident_design.md融合算子状态 buffer 跨循环常驻 双缓冲设计建议【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考