ascend-transformer-boost 中 SortOperation 深度解析：TopK 降序排序与索引输出的双 Runner 实现

ascend-transformer-boost 中 SortOperation 深度解析TopK 降序排序与索引输出的双 Runner 实现【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boostSortOperation 是 CANN ascend-transformer-boost 加速库中用于后处理排序的核心算子它在输入 tensor 的最后一维上做降序排列、仅保留最大的 num 个元素并同时输出排序后的数值与各元素在原输入中的索引。本文以.agent/knowledge知识体系中 sort 条目为骨架结合 路由文件 的源码导航深入 src/ops/ops_infer/sort 与 src/kernels/kernels/sort 的实现细节完整梳理参数约束、InferShape 逻辑、OpsRunner 与 ACLNNRunner 双执行路径、Kernel 内核选择以及测试对齐方法帮助你快速掌握该算子的调用方式与底层原理。1. 算子定位后处理场景下的降序 TopKSortOperation 在知识条目中被归类为category: other、type: single的单输入算子语义与 PyTorch 的torch.topk(..., largestTrue)一致。其功能在参数头文件中有明确描述实现输入 tensor 在最后一维上降序排列并保留最大的 num 个元素输出排序后的 tensor 及各元素对应的索引。—— include/atb/infer_op_params.h典型应用是推理后处理例如在 logits 的最后一维上取前 k 个最大值的数值及其位置用于后续的采样、重排或检索。算子整体数据流为输入 x (shape: [..., last_dim]) ──► SortOperation ──► 输出 values (shape: [..., num]) └─► 输出 indices (shape: [..., num], dtype INT32)从源码结构看该算子同时存在两条执行路径详见第 5 节默认平台走原生 Ops 内核路径SortOpsRunner而 Ascend 950 平台走 ACLNN 封装路径SortAclnnRunner二者共用同一套参数与 InferShape 逻辑。2. 源码文件清单与推荐阅读路线知识条目与路由文件给出了该算子完整的源码构成6 个文件全部位于 src/ops/ops_infer/sort#文件角色1sort_aclnn_runner.cppACLNN RunnerAscend 950 平台2sort_aclnn_runner.hACLNN Runner 头文件3sort_operation.cppOperation 定义InferShape / CreateRunner 决策4sort_operation.hOperation 定义头文件5sort_ops_runner.cpp原生 Ops Runner6sort_ops_runner.h原生 Ops Runner 头文件路由文件推荐的阅读顺序如下可帮助由表及里理解该算子顺序文件重点关注1sort_operation.h输入输出数量、InferShape 签名2sort_operation.cppCreateRunner()的平台决策逻辑3sort_aclnn_runner.hACLNN API 封装接口与函数指针4sort_aclnn_runner.cppWorkspace 计算 ACLNN API 调用链5sort_ops_runner.h原生 Ops 执行接口6sort_ops_runner.cpp原生 Ops 调用链 平台适配相关的 Kernel 层源码位于 src/kernels/kernels/sort参数定义位于 include/atb/infer_op_params.h主索引见 .agent/knowledge/README.md。3. 参数定义与合法性约束SortOperation 使用结构体infer::SortParam定义于 include/atb/infer_op_params.hstruct SortParam { //! 排序后保留的最大的元素的数量 //! \warning num 是一个仅含有一个值的 SVector该值需大于 0 且小于等于输入 x 最后一维的大小 SVectorint32_t num; uint8_t rsv[8] {0}; // 预留参数 }; inline bool operator(const SortParam left, const SortParam right) { return left.num right.num; // 参数比较仅关注 num }关键参数说明numSVectorint32_t仅允许包含 1 个值表示降序后保留的最大元素个数 k等价于 TopK 的 k。合法范围是0 num[0] 输入 x 最后一维的大小。rsv预留字段固定为 8 字节 0用于后续参数扩展保持 ABI 稳定。operator仅比较num用于算子参数缓存与复用时的快速比对。ParamCheckImpl在 sort_operation.cpp 中完成双重校验通用约束num.size() 1且0 num[0] 输入最后一维大小否则返回ERROR_INVALID_PARAMAtlas 300I Duo 推理产品310P特殊约束当num[0] 16时除最后一维外的所有维度均须满足dim LIMITED_NUMLIMITED_NUM 580000000否则返回ERROR_INVALID_TENSOR_DIM——这是为了规避小 k 场景下超大 batch 维度带来的资源开销。此外在 Ascend 950 平台上ParamCheck还会要求num非空且大于 0并触发SortAclnnRunner::LoadMethod()动态加载 ACLNN 函数sort_operation.cpp。4. Operation 层InferShape 与 Runner 路由决策SortOperation继承自OperationBase在 sort_operation.cpp 中定义了输入输出数量与形状推导static const uint32_t IN_TENSOR_NUM 1; // 1 个输入 x static const uint32_t OUT_TENSOR_NUM 2; // 2 个输出 values indices Status SortOperation::InferShapeImpl(const SVectorTensorDesc inTensorDescs, SVectorTensorDesc outTensorDescs) const { outTensorDescs.at(0) inTensorDescs.at(0); outTensorDescs.at(1) inTensorDescs.at(0); outTensorDescs.at(0).shape.dims[outTensorDescs.at(0).shape.dimNum - 1] param_.num[0]; outTensorDescs.at(1).shape.dims[outTensorDescs.at(0).shape.dimNum - 1] param_.num[0]; outTensorDescs.at(1).dtype ACL_INT32; // 索引输出固定为 INT32 return NO_ERROR; }推导规则清晰明了两个输出均继承输入的 shape仅将最后一维替换为num[0]输出 0valuesdtype 与输入一致FP16/BF16/FP32输出 1indicesdtype固定为ACL_INT32与 PyTorchtopk返回的int64索引不同使用时需注意类型差异。Runner 的选择逻辑在CreateRunner中完成sort_operation.cppstd::shared_ptrRunner SortOperation::CreateRunner(Context context) const { if (Mki::PlatformInfo::Instance().GetPlatformType() Mki::PlatformType::ASCEND_950) { return std::make_sharedSortAclnnRunner(param_); } return std::make_sharedSortOpsRunner(param_); }即Ascend 950 → ACLNN Runner其余平台 → 原生 Ops Runner。同时构造函数会根据平台选择不同的算子 IR 配置910B 与 950 使用SortOperationA2and950其余平台使用SortOperationsort_operation.cppIR 配置决定了算子的 tiling/内存管理等元信息。SetParam在更新参数后会清空runner_以强制重建执行器。5. 双 Runner 执行路径详解5.1 SortOpsRunner原生 Ops 内核路径默认平台sort_ops_runner.cpp 通过构建kernelGraph_将算子描述下发到 MKI 内核框架SortOpsRunner::SortOpsRunner(const infer::SortParam param) : OpsRunner(SortOpsRunner), param_(param) { kernelGraph_.inTensors.resize(1); // 输入 x kernelGraph_.outTensors.resize(OUT_TENSOR_COUNT); // values indices ... Mki::SVectorint32_t num; TensorUtil::AtbSVector2OpsSVector(param_.num, num); AsdOps::OpParam::Sort sortNodeParam {num}; // 参数透传给内核层 sortNode.opDesc {0, SortOperation, sortNodeParam}; sortNode.inTensors {xTensor}; sortNode.outTensors {outTensor, indices}; }要点ATB 的SortParam.num通过TensorUtil::AtbSVector2OpsSVector转换为内核层的AsdOps::OpParam::Sort内核图节点名为SortOperation携带 1 输入 2 输出文件末尾的REG_RUNNER_TYPE(SortOpsRunner)与REG_OP_PARAM(AsdOps::OpParam::Sort)完成 Runner 类型与参数类型的注册sort_ops_runner.cpp。5.2 SortAclnnRunnerACLNN 封装路径Ascend 950Ascend 950 平台不直接跑库内自研内核而是复用 CANN ACLNN 的 TopK 算子再叠加一次类型转换。核心实现位于 sort_aclnn_runner.cpp分四步1函数指针动态加载LoadMethodL205-L223通过LoadFromSharedObjectFile从共享库加载两组 API——aclnnTopkGetWorkspaceSize/aclnnTopk与aclnnCastGetWorkspaceSize/aclnnCast并以静态函数指针缓存避免重复加载。2构建 VariantPackBuildAclnnVariantPackL53-L123将 ATB 的输入/输出 Tensor 通过CallAclCreateTensor转为aclTensor由于aclnnTopk返回的索引是 INT64而 SortOperation 输出要求 INT32此处额外创建一个indices_临时 tensorACL_INT64作为 TopK 与 Cast 之间的中转缓冲其缓冲区大小为indicesBufferSize_ outTensor.dataSize * 2。3Workspace 计算与 Executor 创建SetAclNNWorkspaceExecutorL125-L165以k num[0]、dim -1、largest true、sorted true调用 TopK 的 Workspace 接口再对indices_调用 Cast目标ACL_INT32的 Workspace 接口最终 workspace 总量为三者之和this-atbVariantPack_.workspaceBufferSize this-topkWorkspaceSize_ this-castWorkspaceSize_ this-indicesBufferSize_;4流式执行LaunchAclnnKernelL167-L202在executeStream上依次执行aclnnTopkExecute与aclnnCastExecute其中临时indices_通过aclSetOutputTensorAddr/aclSetInputTensorAddr精确指向 workspace 中topkWorkspaceSize_ castWorkspaceSize_偏移处实现内存零额外申请。两个 executor 的IsRepeatable()组合决定执行器是否可复用executorRepeatable_从而支持参数不变时的 kernel 缓存加速。6. Kernel 层TopKDesc 内核选择与 Tiling原生路径最终落到 src/kernels/kernels/sort/sort_operation.cpp 中的内核层SortOperation它根据输出 dtype 选择具体 kernelKernel *GetBestKernel(const LaunchParam launchParam) const override { auto dtype launchParam.GetOutTensor(0).desc.dtype; if (dtype TENSOR_DTYPE_FLOAT16) { return GetKernelByName(TopKDescF16Kernel); } else if (dtype TENSOR_DTYPE_BF16) { return GetKernelByName(TopKDescBF16Kernel); } else if (dtype TENSOR_DTYPE_FLOAT) { return GetKernelByName(TopKDescF32Kernel); } else { return nullptr; // 仅支持 FP16 / BF16 / FP32 } }可见内核层仅支持三种浮点类型FP16TopKDescF16Kernel、BF16TopKDescBF16Kernel、FP32TopKDescF32Kernel对应的 kernel 实现在 src/kernels/kernels/sort/top_k/top_k_desc_kernel.cpp。内核层的 InferShape 与 ATB 层保持一致输出 0 继承输入 shape 并将最后一维改为num[0]输出 1 为TENSOR_DTYPE_INT32、TENSOR_FORMAT_NDsrc/kernels/kernels/sort/sort_operation.cpp。Tiling 阶段位于 src/kernels/kernels/sort/tiling/sort_tiling.cpp将num作为TENSOR_DTYPE_INT32的常量输入随同输入/输出描述一并下发给 runner驱动 kernel 按 batch 维度切分并在每行的最后一维执行 TopK。7. 测试验证与 torch.topk 的 Golden 对齐仓库为该算子提供了完整的算子级与内核级测试用于校验实现的正确性tests/apitest/opstest/python/operations/sort/test_sort.py算子级测试golden 直接使用torch.topk(in_tensors[0], knum[0], largestTrue)并取indices.int()与输出 1 对齐。测试覆盖(10, 22, 4096)的 3D FP16 输入num 3000与num 1500两个场景含参数更新execute_update_param场景同时显式跳过Ascend910A/Ascend310B两个不支持的平台。tests/apitest/kernelstest/sort/test_topk_fp32.py内核级测试op_test.only_910b限定 910B 平台输入为(8, 6)FP32num [1]。值得关注的是其比对策略# Indices (out_tensors[1]) are unstable due to unstable sorting. index_gather torch.gather(self.in_tensors[0], -1, out_tensors[1].to(torch.int64)) result1 torch.equal(index_gather, golden_out_tensors[0])由于排序算法对相等元素的索引顺序不保证稳定测试不直接比对 indices 与 golden 相等而是用输出的索引回查输入、再与 golden 的 values 比较——这一策略既容忍了不稳定排序又严格校验了索引-数值的一致性是验证 TopK 类算子时的最佳实践。tests/apitest/kernelstest/sort/test_topk_bf16.pyBF16 内核测试与 FP32 用例结构一致。8. 使用要点与平台适配小结综合以上源码分析使用 SortOperation 时需注意以下几点参数语义num只接受单元素SVectorint32_t取值须满足0 num[0] 输入最后一维大小否则InferShapeCheck/SetupCheck会返回参数错误输出形状values与indices形状相同输入形状 最后一维截为num[0]indices固定为 INT32与torch.topk默认返回的 INT64 不同若直接对接 PyTorch 侧逻辑需自行转换支持的数据类型内核路径仅支持 FP16 / BF16 / FP32索引稳定排序不做保证涉及相等元素时应采用「索引回查输入」的校验/使用方式平台差异Ascend 950 走aclnnTopk aclnnCast的 ACLNN 路径额外产生一次 INT64→INT32 转换910B 及默认平台走库内TopKDesc*原生内核310PAtlas 300I Duo推理产品对非最后一维有 5.8 亿的上限约束910A / 310B 明确不支持动态参数通过SetParam更新num后 runner 会被置空重建因此同一算子实例支持在推理过程中动态调整 k 值测试中的execute_update_param场景即验证了这一点。如需进一步深入可按第 2 节的阅读顺序从 sort_operation.cpp 读起逐层跟进 sort_ops_runner.cpp原生内核图与 sort_aclnn_runner.cppACLNN 封装两条路径并结合 sort_tiling.cpp 与三个测试文件验证自己的理解。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考