
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.mdop: {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 为最高级知识条目已 completetype: single—— 单输入单输出算子对应源码中GetInputNum() 1、GetOutputNum() 1Runner: ops_runner / aclnn_runner—— 该算子同时具备两条执行路径原生 Ops Runner 与基于 ACLNN 的 RunnerPipeline: 单阶段—— 无需多阶段图编排一次运行即可完成计算。路由文件 .agent/knowledge/routing/cumsum.md 进一步给出分类为infer、复杂度M、共6 个源文件并列出推荐阅读顺序先看 Operation 定义再看 Runner 与 Kernel。二、参数定义与约束CumsumParam算子参数定义在 include/atb/infer_op_params.hinfer::CumsumParam其语义与取值范围如下字段类型默认值说明axesSVectorint64_t无指定计算累加和的轴维度只能包含一个轴索引且必须小于输入张量的维度数exclusiveboolfalse累加结果是否从第 0 个元素开始。false包含首元素y[0] x[0]true排他式前缀和当前暂不支持reverseboolfalse是否反向累加。false正向前缀和true从末尾反向累加当前暂不支持rsvuint8_t[14]{0}预留字段保持参数结构体大小稳定参数创建时的校验逻辑位于 src/ops/ops_infer/cumsum/cumsum_operation.cpp 的CreateOperationinfer::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.hOperation 定义输入输出数量、InferShape 签名cumsum_operation.cppOperation 实现参数校验 CreateRunner()决策cumsum_ops_runner.h原生 Ops Runner 执行接口cumsum_ops_runner.cpp原生 Ops 调用链Fill Cumsum 组合图cumsum_aclnn_runner.hACLNN API 封装接口两段式接口函数指针cumsum_aclnn_runner.cppWorkspace 计算 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.cppInferShapeCheck / SetupCheck校验axes[0]是否落在[0, dimNum)区间内越界即报ERROR_INVALID_PARAMcumsum_operation.cppCreateRunner当前实现直接创建CumsumOpsRunnerstd::make_sharedCumsumOpsRunner(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 OpsRunnerFill Cumsum 的组合内核图CumsumOpsRunner 在构造时构建了一个包含2 个节点的 KernelGraphkernelGraph_.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 并创建 executoraclnnExecuteFunc_签名(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继承自KernelBaseCanSupport校验参数类型为OpParam::Cumsum且输入输出均为 1 个 tensorGetTilingSize计算 tiling buffer 大小与轴常量张量int64_t之和InitImpl调用CumsumTiling()完成 tiling将输出 tensor 索引TENSOR_CUMSUM_OUTPUT_IDX 2设置为Memset 目标SetMemsetInfo即先清零输出并把轴索引TENSOR_CUMSUM_AXIS_IDX 1作为常量张量写入 kernelInfo。针对不同数据类型派生了多个内核并注册内核类支持 dtype用途CumsumF16KernelTENSOR_DTYPE_FLOAT16FP16 前缀和CumsumBF16KernelTENSOR_DTYPE_BF16BF16 前缀和CumsumF16DtmKernelFP16deterministic确定性变体CumsumBF16DtmKernelBF16deterministic确定性变体内核选择逻辑在 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.cumsumdef hillis_steele_fp16(x, axis): # CumsumF16Kernel 的 fp16 累加实现是 Hillis-Steele 并行前缀和 # golden 必须逐位模拟该算法实测与算子输出 maxAbs0 # 用顺序 np.cumsum 或 fp32 累加都会与算子输出产生 fp16 ULP 级偏差。 x np.array(x, dtypenp.float16, copyTrue) 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/910atest_cumsum_dtm(128, 1111){axis: [1], deterministic: True}跳过 310b/910atest_cumsum_bf16(128, 2222){axis: [1]}仅 910btest_cumsum_bf16_dtm(128, 1111){axis: [1], deterministic: True}仅 910b容差上BF16 用rtol/atol 2^-7FP16 用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提供 FillCumsum 组合内核图与aclnnCumsumV2两段式封装两种路径内核层cumsum_kernel.cpp cumsum_tiling.cppFP16/BF16 两类内核、确定性变体、TBE Tiling验证层test_cumsum.pyHillis-Steele 逐位模拟 golden平台差异化覆盖。对于需要在 Ascend 上做序列前缀和、位置偏移累计等后处理计算的开发者这份从知识条目到内核的完整链路即是可直接参考的实战蓝本。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考