asc-devkit 中基于 SoftmaxFlashV2 高阶 API 实现 FlashAttention-2 增量 Softmax 单算子

本文以 asc-devkit 开源仓库中的 softmaxflashv2 样例 为核心，系统讲解如何使用 AscendC::SoftmaxFlashV2 高阶 API 在大语言模型（LLM）注意力机制场景下实现 softmaxflashv2 单算子。该 API 是 SoftmaxFlash 的增强版本，对应 FlashAttention-2 算法，支持 update 模式（增量计算），常用于注意力机制计算与长序列分块处理。阅读本文后，你将掌握该高阶 API 的数学语义、Tiling 与分核设计思路、kernel 实现框架，以及完整的编译、运行与精度验证方法。

SoftmaxFlashV2 的定位与应用场景

在 LLM 的注意力机制（Attention）中，Softmax 计算是最核心也最容易成为性能瓶颈的环节之一。FlashAttention 系列算法通过"分块 + 增量更新"的方式避免把完整的注意力矩阵驻留在片上，从而在长序列场景下显著降低访存开销。其中 FlashAttention-2 进一步优化了并行策略，而 SoftmaxFlashV2 正是 asc-devkit 中与该算法配套的 Softmax 增强版高阶 API。

与基础 Softmax 相比，SoftmaxFlashV2 有两个关键能力（见样例源码注释 softmaxflashv2.asc）：

• 支持 update 模式（增量计算）：可以在已有 max/sum 统计结果的基础上继续累加新分块，避免重复扫描；
• 支持 K 方向（列方向）分块：即沿输入列方向切分并逐块更新统计量，适配长序列场景。

该 API 对应的数学语义定义在接口头文件 softmaxflashv2.h 中，可直接对应 FlashAttention-2 论文《FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning》中描述的增量 Softmax 计算流程。

数学语义：update 模式与非 update 模式

样例将输入 tensor [m0, m1, ..., mt, n]（t ≥ 0）的非尾轴长度相乘的结果视为 m，即把输入 shape 等价为 [m, n]，对每一行独立进行如下计算。其中 x、inmax、insum 为输入，M、S、E 均为输出。

update 为 false（首次计算，不复用历史统计量）：

$$M = rowmax(x_i)$$

$$SoftmaxFlashV2(z_i) = exp(x_i - M)$$

$$S=\sum_{i}^{n} \exp(x_i - M)$$

update 为 true（增量计算，复用之前的 max/sum 结果）：

$$M = max(rowmax(x_i), inmax)$$

$$SoftmaxFlashV2(z_i) = exp(x_i - M)$$

$$E = exp(inmax_i - M)$$

$$S = \sum_{i}^{n} exp(x_i - M) + E \cdot insum$$

可以看到：update 模式下，新的 M 取当前块行最大值与历史 inmax 的较大者，新的 S 在重新归一化当前块的同时，用修正因子 E = exp(inmax - M) 把历史 insum 折算进来。这正是 FlashAttention 系列算法在 K 方向分块累加时的标准"rescaling"技巧：保证无论分多少块，最终统计结果与一次性计算完全一致。

样例规格

本样例为固定 shape 实现，输入与输出的规格如下表所示（摘自 README_en.md）：

项目	name	shape	数据类型	format
样例类型（OpType）	Softmaxflash	-	-	-
样例输入	x	[960, 960]	float	ND
样例输出	max	[960, 8]	float	ND
样例输出	sum	[960, 8]	float	ND
核函数名	softmaxflashv2_custom	-	-	-

输出 shape 为 <a href="https://link.gitcode.com/i/404718afc7a75d8a2cfab29954bbefc1" target="_blank">960, 8] 的原因在于：max 与 sum 是按行统计的归约结果，每行输出 8 个 float 元素（对应一个 32 字节 Block 的元素数，即 FLOAT_NUM_OF_SINGEL_BLOCK = 8，见 [softmaxflashv2.asc），与 impl 层 SoftMaxFlashV2TilingFuncImpl 中 elementNumPerBlk = ONE_BLK_SIZE / dataTypeSize2 的计算逻辑一致（softmax_flashv2_base_impl.h）。

目录结构与数据流

样例目录结构如下：

├── softmaxflashv2
│   ├── scripts
│   │   ├── gen_data.py         // 输入数据和真值数据生成脚本
│   ├── CMakeLists.txt          // 编译工程文件
│   ├── data_utils.h            // 数据读入写出函数
│   ├── softmaxflashv2.asc      // Ascend C 算子实现 & 调用样例
│   └── README.md               // 样例说明文档

运行期数据流为：

1. gen_data.py 生成随机输入 input_x.bin（np.random.uniform(-1, 1, (960, 960))，float32），并用 NumPy 按 Softmax 公式计算真值 golden_max.bin、golden_sum.bin；
2. Host 端通过 ACL 接口读入输入、分配内存、调用内核（<<<>>> 内核调用符），产出 output_max.bin、output_sum.bin；
3. softmaxflashv2.asc 中的 CompareResult 将输出 sum 与真值 sum 逐元素对比（绝对误差与相对误差均需小于 EPS = 1e-5），通过则打印 test pass!。

Kernel 实现：调用高阶 API 完成核心计算

Kernel 侧核心逻辑位于 MyCustomKernel::KernelSoftmax 类（softmaxflashv2.asc），整体遵循 CopyIn → Compute → CopyOut 的流水线结构，并通过 TQue 双缓冲（BUFFER_NUM = 2）隐藏搬运与计算开销。

计算路径：M 方向分核 + K 方向分块

Process 的循环结构（softmaxflashv2.asc）为：

• 外层按 M 方向（行）的 singleCoreLoopCount 循环；
• 内层按 K 方向（列）的 loopK 循环，每个分块大小为 splitK；
• K 方向尾块（tailK > 0）单独处理；
• 每轮循环结束后通过 SetFlag/WaitFlag 同步 V_MTE3 事件，再把该轮 max、sum 统计结果拷出。

调用高阶 API 完成 SoftmaxFlashV2 计算

Compute 中根据是否首块（columnIndex == 0）选择 update 模式，并依据对齐条件选择 BasicBlock 高性能路径（softmaxflashv2.asc）：

AscendC::SoftMaxShapeInfo srcShape = {rowNum, columnNum, rowNum, columnNum};
bool enableBasicBlock =
    (rowNum % BASIC_BLOCK_ROW_FACTOR == 0 && columnNum % BASIC_BLOCK_COLUMN_FACTOR == 0 &&
     columnNum < BASIC_BLOCK_MAX_COLUMN_LENGTH);   // 行对齐 8、列对齐 64、列长 < 2048

if (columnIndex == 0) {   // isUpdate = false：首次计算
    AscendC::SoftmaxFlashV2<float, false, true, true>(
        xLocal, sumLocal, maxLocal, xLocal, expmaxLocal, sumLocal, maxLocal, tmpBuffer, softmaxTiling, srcShape);
} else {                  // isUpdate = true：增量计算
    AscendC::SoftmaxFlashV2<float, true, true, true>(
        xLocal, sumLocal, maxLocal, xLocal, expmaxLocal, sumLocal, maxLocal, tmpBuffer, softmaxTiling, srcShape);
}

其中模板参数依次为：<T, isUpdate, isReuseSource, isBasicBlock>，与接口定义 softmaxflashv2.h 的签名一一对应。srcShape 中的 srcM/srcK 与 oriSrcM/oriSrcK 均设为当前 tile 的行列数，便于 impl 层校验 tiling 与输入 shape 是否匹配（不匹配时 API 内部会依据 workspace 大小重新推导 tiling，见 softmax_flashv2_base_impl.h）。

高阶 API 的两种调用形态

从接口头文件看，AscendC::SoftmaxFlashV2 提供了多种重载（softmaxflashv2.h），核心区别在于：

• 临时工作区形态：最后一个参数既可以是不带 sharedTmpBuffer 的版本（impl 内部通过 PopStackBuffer 从栈上申请 float 工作区，见 softmax_flashv2_base_impl.h），也可以是显式传入的 LocalTensor<uint8_t>& sharedTmpBuffer。样例采用的是后者，其大小由 Tiling 阶段通过 GetSoftMaxFlashV2MinTmpSize 计算得出；
• 数据类型形态：既支持同型输入输出（T 与 T），也支持 half 输入 + float 统计量（half 与 float）的混合形态；
• 输出形态：除标准输出外，还提供带 outReduceMax 的扩展版本（SoftmaxFlashV2MaxImpl）。

Tiling 实现：分核、分块与 API Tiling 参数生成

Tiling 逻辑位于 SoftmaxflashCustomTiling::ComputeTiling（softmaxflashv2.asc），流程如下：

1. 平均分配法分核：先按核数向上对齐总行数（alignedRowNum = (rowNum + coreNum - 1) / coreNum * coreNum），得到每个 Core 的计算行数 coreRowNum，再求出尾核行数 tailCoreRowNum 与实际使用核数 usedNumBlocks；
2. 计算主核内部循环分块参数：调用 GetSingleCoreLoopParam 依据列数选择分块策略（softmaxflashv2.asc）：

列数范围	splitK（K 方向分块）	每次循环行数
colNum ≥ 1024	1024	8
512 ≤ colNum < 1024	512	16
128 ≤ colNum < 512	128	32
colNum < 128	colNum（不分块）	64

3. 计算尾核参数：尾核行数不为 0 时，按尾核行数重新取循环分块参数；
4. 获取临时缓冲区大小并生成 API 所需 Tiling：以主核单次循环的计算 shape 调用 AscendC::GetSoftMaxFlashV2MinTmpSize 获取最小临时空间（字节），再用 AscendC::SoftMaxFlashV2TilingFunc 生成 SoftMaxTiling（此处 localworkspaceSize 同时写入 tiling 结构供 kernel 端初始化 sharedTmpBuffer）；
5. 填充自定义 tiling 结构：SoftmaxflashCustomTilingData 中除 API 内部所需的 SoftMaxTiling softmaxTilingData 外，还包含 columnLength/rowLength/splitK/loopK/tailK 以及主核/尾核各自的循环分块参数。

值得说明的是：样例中只有主核的计算 shape 调用高阶 API 的 tiling 函数；尾核所需的 API tiling 参数由 kernel 侧自行处理——从代码看，尾核复用与主核一致的 softmaxTiling，仅通过不同的 singleLoopCoreRowNum 等循环参数驱动计算。

SoftMaxFlashV2TilingFunc 内部的关键推导在 impl 层 softmax_flashv2_base_impl.h：依据输入 shape、数据类型大小与 workspace 大小计算 reduceM/reduceK/splitM/splitK/rangeM/tailM 等字段，用于指导 API 在 kernel 内部按 M 方向分块归约。SoftMaxTiling 与 SoftMaxShapeInfo 等结构定义于 softmax_tilingdata.h 与 softmax_utils.h。

调用实现：内核调用符与 ACL 流程

Host 端（softmaxflashv2.asc）流程为：

1. aclInit → aclrtSetDevice → aclrtCreateContext → aclrtCreateStream 完成 ACL 初始化；
2. 用 aclrtMallocHost / aclrtMalloc 分别分配 Host 与 Device 内存，其中输入 x 为 960×960×4 字节，workspace 为 1024×4 字节，tiling 数据为 32×4 字节；
3. 调用 Host 端 GenerateTiling(ROW_NUM, COLUMN_NUM, USED_CORE_NUM, ...) 生成 tiling 并通过 memcpy_s 填入缓冲区；
4. aclrtMemcpy 将输入、workspace、tiling 拷入 Device；
5. 使用内核调用符启动核函数：

softmaxflashv2_custom<<<USED_CORE_NUM, 0, stream>>>(xDevice, maxDevice, sumDevice, workspaceDevice, tilingDevice);

其中 USED_CORE_NUM = 40 为使用的核数，核函数签名（softmaxflashv2.asc）接收 x/max/sum/workspace/tiling 五个全局内存指针，先 CopyTiling 把 tiling 数据从全局内存拷入局部，再执行 KernelSoftmax::Init 与 Process；

6. aclrtSynchronizeStream 同步后，将输出拷回 Host 并写入 output/output_max.bin、output/output_sum.bin，与真值比对后打印 test pass! 或 test failed!；
7. 释放内存与 ACL 资源收尾。

编译与运行

在样例根目录 softmaxflashv2 下执行如下步骤。

配置环境变量

请根据当前环境上 CANN 开发套件包的安装方式配置环境变量：

source ${install_path}/cann/set_env.sh

说明：${install_path} 为 CANN 包安装目录，未指定安装目录时默认安装至 /usr/local/Ascend 下。

编译并执行（默认 NPU 模式）

mkdir -p build && cd build;
cmake -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # 默认 npu 模式
python3 ../scripts/gen_data.py
./demo

CPU 调试与 NPU 仿真模式

使用 CPU 调试或 NPU 仿真模式时，添加 -DCMAKE_ASC_RUN_MODE=cpu 或 -DCMAKE_ASC_RUN_MODE=sim 参数即可：

cmake -DCMAKE_ASC_RUN_MODE=cpu -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # cpu 调试模式
cmake -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # NPU 仿真模式

注意：切换编译模式前需清理 cmake 缓存，可在 build 目录下删除 CMakeCache.txt 后重新执行 cmake。

编译选项说明

选项	可选值	说明
CMAKE_ASC_RUN_MODE	npu（默认）、cpu、sim	运行模式：NPU 运行、CPU 调试、NPU 仿真
CMAKE_ASC_ARCHITECTURES	dav-2201（默认）、dav-3510	NPU 架构：dav-2201 对应 Atlas A2 系列产品和 Atlas A3 系列产品；dav-3510 对应 Ascend 950PR&950DT 系列产品

上述选项在 CMakeLists.txt 中以 CACHE STRING 形式定义，并最终通过 --npu-arch=${CMAKE_ASC_ARCHITECTURES} 编译选项传给 ASC 编译器（同文件第 34-36 行）。

执行结果

执行结果如下，说明精度对比成功：

test pass!

支持的产品与 CANN 软件版本

产品	CANN 软件版本
Ascend 950PR&950DT 系列产品	>= CANN 9.1.0
Atlas A3 系列产品	>= CANN 9.0.0
Atlas A2 系列产品	>= CANN 9.0.0

从接口实现看，SoftmaxFlashV2 在 softmaxflashv2.h 中通过 __NPU_ARCH__ 宏对 2201/2002/3510/9201/3002/5102/5162/3003/3113 等架构开放，并按架构分别展开为 membase（v220/v200）与 regbase（3510/v300）的底层实现（softmax_flashv2_base_impl.h），与上表支持范围相互印证。

精度验证机制

真值由 gen_data.py 中的 softmax_py_float 用 NumPy 逐通道计算行最大值与行指数和生成。Kernel 输出的 sum 结果在 CompareResult 中与 golden_sum.bin 比对，判定条件为：绝对误差 |a-b| 与相对误差 |a-b|/|b| 均须小于 EPS = 1e-5，任一元素超差即判定失败（softmaxflashv2.asc）。

进一步的源码参考

• 高阶 API 完整接口与重载：include/adv_api/activation/softmaxflashv2.h
• Tiling 查询与生成函数（含 GetSoftMaxFlashV2MinTmpSize/GetSoftMaxFlashV2MaxTmpSize/SoftMaxFlashV2TilingFunc）：include/adv_api/activation/softmax_tiling.h
• 公共数据结构（SoftMaxShapeInfo/SoftmaxConfig/SOFTMAX_DEFAULT_CFG）：include/adv_api/activation/softmax_utils.h
• impl 层核心实现（tiling 推导与计算入口）：impl/adv_api/detail/activation/softmax/softmax_flashv2_base_impl.h
• 单元测试（覆盖 isUpdate/isBasicBlock/混合精度等组合）：tests/api/adv_api/activation/softmax/test_operator_softmax_v220.cpp
• API 参数检查用例：tests/api/adv_api/api_check/activation/softmax/softmax_flashv2/kernel_softmax_flashv2_check.cpp

整体来看，softmaxflashv2 样例是理解"高阶 API 快速实现 + 自定义 Tiling 精细化控制"这一典型 Ascend C 开发范式的最小完整示例：用户只需把数据搬入 UB，调用一行 SoftmaxFlashV2 高阶 API 即可完成 FlashAttention-2 风格的增量 Softmax 计算，再配合按行分核、按列分块的 Tiling 策略，即可在长序列注意力场景中获得可扩展的多核并行能力。