ATB Cumsum 算子深度解析：前缀和在 Ascend Transformer Boost 中的实现与调用链
【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost
导读
Cumsum（累计和/前缀和）是 Transformer 推理与训练后处理中常用的基础算子，例如用于位置偏移换算、序列长度前缀和等场景。本文以 ascend-transformer-boost 仓库中的知识条目 .agent/knowledge/ops/other/cumsum/index.md 为骨架，结合路由文件与src/ops/ops_infer/cumsum/、src/kernels/kernels/cumsum/的真实源码，完整讲解该算子的参数定义、合法性校验、双 Runner 执行路径（OpsRunner 与 ACLNN Runner）、Kernel 层 FP16/BF16 实现以及精度测试的 Hillis-Steele 模拟方法，帮助读者掌握在 ATB 框架下从图构建到内核执行的全链路工作原理。
一、算子定位：知识条目中的元信息
在 ATB 仓库的 Agent 知识体系中，Cumsum 被登记为一条完整的知识条目，其元数据如下（见 .agent/knowledge/ops/other/cumsum/index.md）：
op: {name: "cumsum", category: "other", tier: "S", type: "single"} source: {repo_path: "src/ops/ops_infer/cumsum/"} knowledge: {status: "complete", last_extracted: "2026-07-06"}其中：
- category: other—— 算子不属于通信、Attention 或 Norm 等专项分类，而是归入通用基础算子集合（与 concat、slice、fill、transpose 等同目录，见 .agent/knowledge/ops/other/）；
- tier: S—— 表示该算子的知识完整度分级（S 为最高级，知识条目已 complete）；
- type: single—— 单输入单输出算子，对应源码中
GetInputNum() == 1、GetOutputNum() == 1； - Runner: ops_runner / aclnn_runner—— 该算子同时具备两条执行路径：原生 Ops Runner 与基于 ACLNN 的 Runner；
- Pipeline: 单阶段—— 无需多阶段图编排，一次运行即可完成计算。
路由文件 .agent/knowledge/routing/cumsum.md 进一步给出：分类为infer、复杂度M、共6 个源文件，并列出推荐阅读顺序（先看 Operation 定义，再看 Runner 与 Kernel）。
二、参数定义与约束：CumsumParam
算子参数定义在 include/atb/infer_op_params.h（infer::CumsumParam），其语义与取值范围如下：
| 字段 | 类型 | 默认值 | 说明 |
|---|---|---|---|
axes | SVector<int64_t> | 无 | 指定计算累加和的轴（维度），只能包含一个轴索引，且必须小于输入张量的维度数 |
exclusive | bool | false | 累加结果是否从第 0 个元素开始。false：包含首元素（y[0] = x[0]）；true：排他式前缀和（当前暂不支持） |
reverse | bool | false | 是否反向累加。false：正向前缀和；true：从末尾反向累加（当前暂不支持） |
rsv | uint8_t[14] | {0} | 预留字段，保持参数结构体大小稳定 |
参数创建时的校验逻辑位于 src/ops/ops_infer/cumsum/cumsum_operation.cpp 的CreateOperation<infer::CumsumParam>模板特化中，任何一项不满足都会直接返回ERROR_INVALID_PARAM并打印错误日志：
OP_PARAM_RSV_CHECK(opParam); if (opParam.axes.size() != 1) { ATB_LOG(ERROR) << "cumsumParam axes size must be 1"; return ERROR_INVALID_PARAM; } if (opParam.reverse) { ... } // "cumsum does not support reverse yet" if (opParam.exclusive) { ... } // "cumsum does not support exclusive yet"也就是说，当前版本对外仅开放单轴、正向、包含首元素的累加能力，exclusive与reverse虽然参数结构体中已预留，但传入true会直接报错。
三、源码文件清单与阅读路径
路由文件列出了 Cumsum 的全部 6 个源文件，它们构成Operation → Runner → Kernel三层结构：
| 文件 | 角色 |
|---|---|
| cumsum_operation.h | Operation 定义：输入输出数量、InferShape 签名 |
| cumsum_operation.cpp | Operation 实现：参数校验 +CreateRunner()决策 |
| cumsum_ops_runner.h | 原生 Ops Runner 执行接口 |
| cumsum_ops_runner.cpp | 原生 Ops 调用链：Fill + Cumsum 组合图 |
| cumsum_aclnn_runner.h | ACLNN API 封装接口（两段式接口函数指针） |
| cumsum_aclnn_runner.cpp | Workspace 计算 +aclnnCumsumV2调用 |
建议按路由文件的顺序阅读：先看cumsum_operation.h/.cpp理解算子对外契约，再分别跟进两条 Runner 路径，最后进入 Kernel 目录。
四、Operation 层：InferShape 与 Runner 决策
CumsumOperation继承自OperationBase（见 src/ops/ops_infer/cumsum/cumsum_operation.h），核心行为：
- 输入/输出数量：各 1 个；
- InferShape：输出张量 desc 直接拷贝输入 desc，即输出 shape、dtype、format 与输入完全一致（
outTensorDescs.at(0) = inTensorDescs.at(0)，见 cumsum_operation.cpp）； - InferShapeCheck / SetupCheck：校验
axes[0]是否落在[0, dimNum)区间内，越界即报ERROR_INVALID_PARAM（cumsum_operation.cpp）； - CreateRunner：当前实现直接创建
CumsumOpsRunner（std::make_shared<CumsumOpsRunner>(param_)，cumsum_operation.cpp）。
值得注意的是，源码中同时存在并注册了CumsumOpsRunner与CumsumAclnnRunner两个 Runner（分别通过REG_RUNNER_TYPE注册，见 cumsum_ops_runner.cpp 与 cumsum_aclnn_runner.cpp），这与知识条目中“Runner: ops_runner / aclnn_runner”的记载一致——ACLNN 路径为算子提供了基于 CANN 统一算子的备选/并行实现。
五、Runner 层：两条执行路径
5.1 OpsRunner：Fill + Cumsum 的组合内核图
CumsumOpsRunner 在构造时构建了一个包含2 个节点的 KernelGraph：
kernelGraph_.nodes.resize(2); // nodes: 2 auto &fillNode = kernelGraph_.nodes[0]; // "FillOperation"：按输入 shape 填充输出 auto &cumsumNode = kernelGraph_.nodes[1]; // "CumsumOperation"：真正的前缀和内核- Fill 节点：
opDesc = {0, "FillOperation", AsdOps::OpParam::Fill()}，通过inferShapePreFunc在 Launch 时依据输入 dims 动态设置Fill({false, {0}, inDim})，负责先把输出张量按 shape 填零（Kernel 层同时设置了MemsetInfo，见下文）； - Cumsum 节点：
opDesc = {0, "CumsumOperation", asdParam}，其中asdParam由 ATB 参数转换而来：
AsdOps::OpParam::Cumsum asdParam; asdParam.exclusive = param_.exclusive; asdParam.reverse = param_.reverse; for (std::size_t i = 0; i < param_.axes.size(); ++i) { asdParam.axis.push_back(param_.axes[i]); } cumsumNode.inTensors = {&xTensor}; cumsumNode.outTensors = {&resultTensor};该 Runner 通过REG_OP_PARAM同时注册了Fill与Cumsum两类算子参数，运行时由 Ops Runner 基础设施按图调度。
5.2 AclnnRunner：两段式 aclnnCumsumV2 封装
CumsumAclnnRunner 采用 CANN ACLNN 标准的两段式接口，通过函数指针在运行时从共享库动态加载：
aclnnGetWorkspaceSizeFunc_：签名(input, dim, exclusive, reverse, output, workspaceSize, executor)，在SetAclNNWorkspaceExecutor()中调用，用于申请 workspace 并创建 executor；aclnnExecuteFunc_：签名(workspace, workspaceSize, executor, stream)，在LaunchAclnnKernel()中真正下发执行。
动态加载发生在LoadMethod()（cumsum_aclnn_runner.cpp），加载的符号为aclnnCumsumV2GetWorkspaceSize/aclnnCumsumV2。关键实现细节：
// 注意：cumsum支持多个轴，这里取第一个轴，因为aclnn接口只支持单个维度 int64_t dim = param_.axes.empty() ? 0 : param_.axes[0]; aclnnStatus ret = aclnnGetWorkspaceSizeFunc_( aclnnVariantPack_.aclInTensors.at(0)->tensor, // 输入tensor dim, // cumsum维度 param_.exclusive, // exclusive参数 param_.reverse, // reverse参数 aclnnVariantPack_.aclOutTensors.at(0)->tensor, // 输出tensor &(atbVariantPack_.workspaceBufferSize), // workspace大小 &raw_executor_ptr); // executor由于 ACLNN 接口只支持单维，源码取axes[0]传入（与参数层强制axes.size() == 1的约束呼应）。BuildAclnnVariantPack()负责把 ATB Tensor 通过CallAclCreateTensor转换为aclTensor，并设置needUpdateTensorDataPtr = true以便复用 tensor 句柄。执行流通过GetExecuteStream(atbVariantPack_.context)获取，保证与 ATB 上下文中的 Stream 对齐。
六、Kernel 层：FP16/BF16 内核与确定性变体
算子内核位于 src/kernels/kernels/cumsum/cumsum_kernel.cpp，核心类CumsumKernel继承自KernelBase：
- CanSupport：校验参数类型为
OpParam::Cumsum，且输入输出均为 1 个 tensor； - GetTilingSize：计算 tiling buffer 大小与轴常量张量（
int64_t）之和； - InitImpl：调用
CumsumTiling()完成 tiling，将输出 tensor 索引TENSOR_CUMSUM_OUTPUT_IDX = 2设置为Memset 目标（SetMemsetInfo，即先清零输出），并把轴索引TENSOR_CUMSUM_AXIS_IDX = 1作为常量张量写入 kernelInfo。
针对不同数据类型派生了多个内核并注册：
| 内核类 | 支持 dtype | 用途 |
|---|---|---|
CumsumF16Kernel | TENSOR_DTYPE_FLOAT16 | FP16 前缀和 |
CumsumBF16Kernel | TENSOR_DTYPE_BF16 | BF16 前缀和 |
CumsumF16DtmKernel | FP16 | deterministic（确定性）变体 |
CumsumBF16DtmKernel | BF16 | deterministic（确定性）变体 |
内核选择逻辑在 src/kernels/kernels/cumsum/cumsum_operation.cpp 的GetBestKernel()中：依据输出 dtype 与param.deterministic标志二选一。若输出 dtype 既不是 FP16 也不是 BF16，则返回nullptr并记录错误日志——当前内核仅支持 FP16 与 BF16 两种精度。
InferShape 层同样校验axis.size() != 1直接失败（"null or multi axis is not support now"），并规定输出与输入保持相同 dtype、format、dims。
Tiling 过程
src/kernels/kernels/cumsum/tiling/cumsum_tiling.cpp 中的CumsumTiling()通过TbeTilingRunner在 Host 侧完成形状与策略计算：
auto runner = AsdOpsGeRt::TbeTilingRunner() .SetName("Cumsum") .SetKernelName(kernelName) .AddInput(tensorDesc0.dtype, tensorDesc0.format, tensorDesc0.dims) .AddConstInput(TENSOR_DTYPE_INT64, TENSOR_FORMAT_ND, axisShape, axis.data(), ...) .AddOutput(tensorDescOut.dtype, tensorDescOut.format, tensorDescOut.dims) .AddAttrBool(param.exclusive) .AddAttrBool(param.reverse); return GetTilingFromRunner(kernelInfo, runner, binHandle);即：输入张量 + 轴常量（ND 格式 int64）+ 输出张量 +exclusive/reverse两个布尔属性一起交给 TBE Tiling Runner，结果回填到kernelInfo供内核加载。
七、精度与测试验证：Hillis-Steele 前缀和
Cumsum 的测试位于 tests/apitest/kernelstest/cumsum/test_cumsum.py，其中对 FP16 路径的 golden 计算做了非常精细的处理——逐位模拟算子的 Hillis-Steele 并行前缀和算法，而不是直接使用np.cumsum：
def hillis_steele_fp16(x, axis): # CumsumF16Kernel 的 fp16 累加实现是 Hillis-Steele 并行前缀和， # golden 必须逐位模拟该算法（实测与算子输出 maxAbs=0）， # 用顺序 np.cumsum 或 fp32 累加都会与算子输出产生 fp16 ULP 级偏差。 x = np.array(x, dtype=np.float16, copy=True) n = x.shape[axis] d = 1 while d < n: cur = [slice(None)] * x.ndim; cur[axis] = slice(d, None) prev = [slice(None)] * x.ndim; prev[axis] = slice(None, -d) x[tuple(cur)] = np.float16(x[tuple(cur)] + x[tuple(prev)]) d *= 2 return x测试覆盖四个用例，并带平台过滤注解（skip_310b、skip_910a、only_910b）：
| 用例 | shape | 参数 | 平台限制 |
|---|---|---|---|
test_cumsum | (128, 2222) | {"axis": [1]} | 跳过 310b/910a |
test_cumsum_dtm | (128, 1111) | {"axis": [1], "deterministic": True} | 跳过 310b/910a |
test_cumsum_bf16 | (128, 2222) | {"axis": [1]} | 仅 910b |
test_cumsum_bf16_dtm | (128, 1111) | {"axis": [1], "deterministic": True} | 仅 910b |
容差上，BF16 用rtol/atol = 2^-7，FP16 用rtol/atol = 2^-8的torch.allclose比较。这个测试文件的价值在于：它揭示了内核实际的累加实现是Hillis-Steele 并行扫描而非顺序扫描，这对任何关心 FP16 前缀和逐位精度的开发者都是第一手证据。
八、如何在 ATB 中配置与使用 Cumsum
结合前述参数约束，在 ATB 图中挂载 Cumsum 的要点如下：
- 参数构造：
infer::CumsumParam中axes必须且只能包含一个轴索引，取值范围[0, 输入维度数)；exclusive与reverse保持默认false； - 输入输出：各 1 个 tensor，输出 shape/dtype/format 与输入完全一致（InferShape 直通），运行时输出张量会先被清零再写入累加结果；
- 精度要求：内核当前仅支持FP16 与 BF16输入输出，其他 dtype 在
GetBestKernel()阶段会失败； - 运行路径：通过
CreateOperation(param, &operation)创建（内部完成上述全部校验），再由框架根据配置选择 OpsRunner 图路径或 ACLNN 路径执行； - 调试辅助：算子支持
ATB_LOG(INFO/DEBUG)日志（如 variantPack 打印、workspaceSize 打印），可在日志与调试指南（见 docs/日志与调试.md）中定位问题。
九、总结
从知识条目到源码，Cumsum 在 ascend-transformer-boost 中呈现出一条清晰的分层实现链路：
- 契约层（infer_op_params.h + cumsum_operation.cpp）：单轴、正向、包含首元素，shape 直通，参数严格校验；
- 执行层（cumsum_ops_runner.cpp 与 cumsum_aclnn_runner.cpp）：提供 Fill+Cumsum 组合内核图与
aclnnCumsumV2两段式封装两种路径； - 内核层（cumsum_kernel.cpp + cumsum_tiling.cpp）：FP16/BF16 两类内核、确定性变体、TBE Tiling；
- 验证层（test_cumsum.py）：Hillis-Steele 逐位模拟 golden，平台差异化覆盖。
对于需要在 Ascend 上做序列前缀和、位置偏移累计等后处理计算的开发者，这份从知识条目到内核的完整链路即是可直接参考的实战蓝本。
【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
