人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载导读本文基于 CANN asc-devkit 开源仓库中的 softmaxflashv3 示例系统讲解如何在大模型训练注意力FlashAttention优化场景下使用昇腾 AscendC 高阶 APIAscendC::SoftmaxFlashV3实现单算子 softmaxflashv3。读完本文你将掌握该 API 的数值计算原理Softmax PASA 算法与 update 两分支公式、SoftMaxParams/SoftMaxTiling关键参数含义、基于平均分配法的多核 Tiling 设计以及从数据生成、编译到精度对比验证的完整工程化流程。背景SoftmaxFlashV3 是什么FlashAttention 类算法会把长序列注意力计算拆分为多个 block 分块迭代每轮迭代都会把上一轮的统计量行最大值、指数和、行均值传入下一轮。SoftmaxFlashV3正是面向这种带历史统计量更新场景设计的高阶 API。从 接口头文件 的注释可以看出它是 SoftmaxFlash 的增强版本对应 Softmax PASAPartitioned Adaptive Softmax Attention算法相比 SoftmaxFlashV2在统计输出中增加了均值mean计算通过alpha 参数对输入进行数值平移shift优化改善数值稳定性适用于对数值精度要求更高的 softmax 计算场景。示例采用half输入/输出与float统计值的数据类型组合输入 Tensor shape 为[8, 2048]完成一次完整的 SoftmaxFlashV3 注意力计算。从 SoftmaxFlashV3Impl 的 static_assert 可以确认该 API 当前支持的数据类型组合为T: half, U: floatU类型用于统计量inMax/inSum/inMean 及输出统计的中间计算。支持的产品与 CANN 软件版本产品CANN 软件版本Ascend 950PR 950DT 系列产品 CANN 9.1.0Atlas A3 系列产品 CANN 9.0.0Atlas A2 系列产品 CANN 9.0.0从接口实现的角度看SoftmaxFlashV3仅在特定 NPU 架构如__NPU_ARCH__为 2201、2002、3510、9201、3002 等下展开AICAI Core 之外的通用计算单元上为空实现这一点在 softmaxflashv3.h 的预处理分支中有明确体现。样例目录结构examples/01_simd_cpp_api/04_advanced_api/01_activation/softmaxflashv3/ ├── scripts │ └── gen_data.py // 输入数据与真值golden数据生成脚本 ├── CMakeLists.txt // 编译工程文件 ├── data_utils.h // 二进制文件读写工具函数 ├── softmaxflashv3.asc // Ascend C 算子实现与调用样例 └── README.md // 样例说明文档中文算法原理与计算公式SoftmaxFlashV3 将输入 Tensor[m0, m1, ..., mt, n]t 0中非尾轴长度m0, m1, ..., mt的乘积视作m即把输入 shape 等价为[m, n]。随后将输入x的尾轴宽度维度切分为splitMeanCnt个块切分后的张量记为x_cnti。下面按update取值给出完整公式x、inmax、insum、inmean为输入M、S、E、A均为输出。update 为 false首轮迭代无历史统计量$$ A_1 \text{rowmean}(x_{cnt})i, i \in [0, \text{splitMeanCnt}]\ A_2 \text{rowmean}(x_i), i \in [0, n]\ x_i x_i - (A_2 - A_1) * (\alpha / (1 - \alpha))\ A A_2\ M_1 \text{rowmax}(x_i), i \in [0, n]\ M M_1\ M_2 M\ \text{SoftmaxFlashV3}(z_i) \exp(x_i - M_2), i \in [0, n]\ S \sum{i}^{n} \exp(x_i - M_2) $$update 为 true迭代更新融合历史统计量$$ A_1 \text{rowmean}(x_{cnt})_i, i \in [0, \text{splitMeanCnt}]\ A_2 \text{rowmean}(x_i), i \in [0, n]\ x_i x_i - (A_2 - A_1) * (\alpha / (1 - \alpha))\ A (A_2 \text{inmean} * (\text{loopCnt} - 1)) / \text{loopCnt}\ M_1 \text{rowmax}(x_i), i \in [0, n]\ C (A_2 - A) * (\alpha / (1 - \alpha))\ P (\text{inmean} - A) * (\alpha / (1 - \alpha))\ M \max(C M_1, P \text{inmax})\ M_2 M - C\ \text{SoftmaxFlashV3}(z_i) \exp(x_i - M_2), i \in [0, n]\ E \exp(\text{inmax}i - M_2 P)\ S \sum{i}^{n} \exp(x_i - M_2) E * \text{insum} $$关键参数splitMeanCnt、alpha 与 loopCnt公式中的三个关键参数在 SoftMaxParams 结构体 中有明确定义示例主函数中传入的具体值见 softmaxflashv3.asc如下参数类型示例取值含义splitMeanCntuint32_t8尾轴切分块数决定A_1 rowmean(x_cnti)的局部均值粒度alphafloat0.9375数值稳定性平移系数公式中以alpha / (1 - alpha)参与缩放loopCntuint32_t2迭代轮数update 分支用(inmean * (loopCnt - 1))加权历史均值srcM/srcKuint32_t8 / 2048输入行数 / 列数oriSrcM/oriSrcKuint32_t8 / 2048原始 shape用于 tiling 校验SoftMaxParams的默认值loopCnt1、splitMeanCnt8、alpha0.9375也定义在 softmax_utils.h 中。alpha 的典型取值 0.9375 15/16是二进制友好的缩放因子。这一计算逻辑与 gen_data.py 中的参考实现alapha 0.9375; scalar alapha / (1 - alapha)完全一致。样例规格项目规格样例类型OpTypesoftmaxflashv3输入srcshape[8, 2048]数据类型half格式ND输入inMaxshape[8, 8]数据类型float格式ND输入inSumshape[8, 8]数据类型float格式ND输入inMeanshape[8, 8]数据类型float格式ND输出dstshape[8, 2048]数据类型half格式ND核函数名softmaxflashv3_custom其中统计输入 inMax/inSum/inMean 的尾维 8 与splitMeanCnt8对应——每行需要 8 个分块的统计值elementNumPerBlk1 32 / sizeof(float) 8elementNumPerBlk2 32 / sizeof(half) 16定义了每个 block 的元素数见 softmaxflashv3.asc这也是统计张量尾维取 8、expMax 中间张量尾维取 16 的由来。关键实现解析Kernel 侧CopyIn → Compute → CopyOut核函数softmaxflashv3_custom从 Global Memory 读取srcGm/inMaxGm/inSumGm/inMeanGm与 tiling 数据实例化KernelSoftmaxFlashhalf, float后依次执行Init与Process见 softmaxflashv3.ascInit根据GetBlockIdx()核 ID计算本核处理数据的偏移量数据块偏移blockLength coreRowNum * width统计块偏移msLength coreRowNum * elementNumPerBlk1并绑定 5 组 Global/Local 缓冲区src、dst、max、sum、mean随后用TPipe::InitBuffer分配队列内存。Process先做分核边界判断超出usedNumBlocks的核直接返回防止无效执行再判定 Tail 核GetBlockIdx() usedNumBlocks且有尾行时以tailCoreRowNum处理剩余行最后串起CopyIn → Compute → CopyOut。Compute通过LocalTensor::SetShapeInfo设置各张量二维 shape 后调用高阶 APIAscendC::SoftMaxParams params {rowNum, width, rowNum, width, loopCnt, splitMeanCnt, alpha}; AscendC::SoftmaxFlashV3T, U, true, false( dstLocal, inmeanLocal, insumLocal, inmaxLocal, srcLocal, expMaxTensor, inmeanLocal, insumLocal, inmaxLocal, tilingData, params);模板参数isUpdatetrue表示启用 update 分支融合历史统计量isReuseSourcefalse为保留参数。API 的完整签名与各模板参数语义见 softmaxflashv3.hisBasicBlock在 shape 满足m % 8 0 k % 64 0时可置 true 以提升性能isDataFormatNZ用于 NZ 格式输入SoftmaxConfig默认关闭 tiling 重算SOFTMAX_DEFAULT_CFG {true, 0, 0, SOFTMAX_NORMAL}。在底层实现中updatetrue路径会依次执行ModifyInputImpl均值平移修正输入、SoftmaxFlashV3UpdateMeanImplmean (inmean * (loopCnt - 1) mean) / loopCnt、行最大值归约、shiftCurr/shiftPrev计算、max(maxTmp shiftCurr, inmax shiftPrev)、exp(inmax - max shiftPrev)、expSum expMax * insum rowsum(y)等一系列向量指令详见 membase/v220/softmax_flashv3_impl.h。同时该实现通过ASCENDC_ASSERT要求params.srcK 512即尾轴长度不得小于 512同文件 L344-L346。Tiling 侧平均分配法多核切分样例的 Tiling 设计分两层Host 侧 Tiling定义SoftmaxFlashV3CustomTilingDataheight、width、elementNumPerBlk1/2、loopCnt、splitMeanCnt、alpha、coreRowNum、tailCoreRowNum、usedNumBlocks 等字段ComputeTiling/GenerateTiling中采用平均分配法alignedRowNum (height coreNum - 1) / coreNum * coreNum先按核数向上取整对齐得到每核行数coreRowNum alignedRowNum / coreNum再计算尾核行数tailCoreRowNum height % coreRowNum与使用的核数usedNumBlocks height / coreRowNum见 softmaxflashv3.asc。Kernel 侧 Tiling核内根据baseM自行计算SoftMaxTiling参数。baseM按片上可用空间推导10 * 1024 / (8 * 5 1024 * 2 64)并要求是 8 的倍数、不超过rowNum当推导结果小于 8 时回退为rowNum要求rowNum为 8 的倍数。随后按last rowNum / baseM、tail rowNum % baseM构造完整的SoftMaxTilingsoftmaxflashv3.asc。SoftMaxTiling结构体字段srcM/srcK/srcSize、outMaxM/outMaxK/outMaxSize、splitM/splitK/splitSize、reduceM/reduceK/reduceSize、rangeM/tailM/tailSplitSize/tailReduceSize在 softmax_tilingdata.h 中定义其中rangeM/tailM描述按splitM切分后的整块行数与尾块行数底层SoftmaxFlashV3NoUpdateExtImpl/SoftmaxFlashV3NDExtImpl正是以rangeM次循环加tailM尾块处理完成整轮计算。Host 侧ACL 运行时与精度验证main函数softmaxflashv3.asc完整演示了单算子调用的标准流程aclInit→aclrtSetDevice→aclrtCreateContext→aclrtCreateStream初始化 ACL 运行时依次为 tiling、4 组输入src/inMax/inSum/inMean、1 组输出分配 Host/Device 内存GenerateTiling生成 tiling 后用aclrtMemcpy拷贝到 Device以内核调用符numBlocks, 0, stream启动核函数softmaxflashv3_custom并aclrtSynchronizeStream同步等待完成结果拷回 Host 写盘调用CompareResult与 golden 对比。精度对比采用绝对误差与相对误差双判据EPS 1e-3任一元素ae EPS re EPS即视为错误softmaxflashv3.asc。golden 数据由 gen_data.py 中的softmax_flash_v3参考算法NumPy 实现is_fp16True、en_updateTrue、输入随机均匀分布[0, 10)逐元素生成写入./output/golden.bin。仓库中还提供了独立的 SoftmaxFlashV3 单元测试其中以loopCnt2、splitMeanCnt8构造SoftMaxParams与SoftMaxTiling并调用同一 API可佐证参数用法的一致性。编译与运行在样例根目录examples/01_simd_cpp_api/04_advanced_api/01_activation/softmaxflashv3/下执行以下步骤。1. 配置环境变量根据当前环境上 CANN 开发套件包的安装方式配置环境变量source ${install_path}/cann/set_env.sh说明${install_path}为 CANN 包安装目录未指定安装目录时默认安装至/usr/local/Ascend下。2. 编译并运行NPU 模式默认mkdir -p build cd build; cmake -DCMAKE_ASC_ARCHITECTURESdav-2201 ..; make -j; # 默认 npu 模式 python3 ../scripts/gen_data.py ./demogen_data.py会在样例目录下生成input/input.bin、inputMax.bin、inputSum.bin、inputMean.bin、tiling.bin与output/golden.bin数据demo程序读取输入执行核函数并完成精度对比。3. CPU 调试 / NPU 仿真模式cmake -DCMAKE_ASC_RUN_MODEcpu -DCMAKE_ASC_ARCHITECTURESdav-2201 ..; make -j; # CPU 调试模式 cmake -DCMAKE_ASC_RUN_MODEsim -DCMAKE_ASC_ARCHITECTURESdav-2201 ..; make -j; # NPU 仿真模式注意切换编译模式前需清理 cmake 缓存可在 build 目录下执行rm CMakeCache.txt后重新 cmake。4. 编译选项说明两个选项均定义在 CMakeLists.txt 中选项可选值说明CMAKE_ASC_RUN_MODEnpu默认、cpu、sim运行模式NPU 运行、CPU 调试、NPU 仿真CMAKE_ASC_ARCHITECTURESdav-2201默认、dav-3510NPU 架构dav-2201对应 Atlas A2 系列与 Atlas A3 系列产品dav-3510对应 Ascend 950PR 950DT 系列产品工程通过find_package(ASC REQUIRED)引入昇腾算子编译工具链并将tiling_api、register、platform、unified_dlog、graph_base等库链接进demo可执行文件同时以--npu-arch${CMAKE_ASC_ARCHITECTURES}指定设备侧编译架构。5. 执行结果精度对比成功时输出test pass!若对比失败程序会输出wrongNum: 数量与test failed!此时可检查gen_data.py生成的输入分布、alpha/splitMeanCnt参数是否与核侧一致或改用 CPU 调试模式定位问题。延伸阅读接口定义include/adv_api/activation/softmaxflashv3.h、include/adv_api/activation/softmax_utils.h底层实现impl/adv_api/detail/activation/softmax/softmax_flashv3_base_impl.h分架构实现见membase/v220、membase/v200、regbase/3510、regbase/v300等目录Tiling 定义include/adv_api/activation/softmax_tilingdata.h、include/adv_api/activation/softmax_tiling.h单元测试tests/api/adv_api/activation/softmax/test_operator_softmaxflashv3_v220.cpp快速入门docs/zh/quick_start.md、docs/en/quick_start.md赞分享人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载相关推荐Ascend C 高阶 API 实践基于 SoftmaxFlashV3 实现 FlashAttention 数值稳定 Softmax 单算子softmaxflashv3 样例深度解析Ascend C 高阶 API 实践基于 SoftmaxFlashV3 实现 FlashAttention 数值稳定 Softmax 单算子softmaxf人工智能深度学习算子库CANNAscendasc-devkit 昇腾 C API基于 SoftmaxFlashV2 高阶 API 实现 FlashAttention-2 注意力 Softmax 算子asc devkit 昇腾 C API基于 SoftmaxFlashV2 高阶 API 实现 FlashAttention 2 注意力 Softmax 算人工智能深度学习算子库CANNAscendasc-devkit SoftmaxFlashV3 样例实战基于 Ascend C 高阶 API 实现 FlashAttention V3 的 Softmax 核心算子asc devkit SoftmaxFlashV3 样例实战基于 Ascend C 高阶 API 实现 FlashAttention V3 的 Softmax人工智能深度学习算子库CANNAscend上一篇如何高效使用chromaticChromium/V8广谱注入的完整专业指南下一篇终极屏幕翻译工具告别复制粘贴实现真正的框选即译创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考