ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

ascend-transformer-boost 中 Unpad 算子的源码路径导航与实现原理

ascend-transformer-boost 中 Unpad 算子的源码路径导航与实现原理 ascend-transformer-boost 中 Unpad 算子的源码路径导航与实现原理【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本文是 ascend-transformer-boost 仓库.agent/knowledge/routing/unpad.md路由文件的深度展开。该路由文件将读者导向 Unpad 算子的四个核心源码文件本文在此基础上结合仓库内 Operation 层、OpsRunner 层、AscendC Kernel 层与测试用例完整还原 Unpad 算子的数据流、InferShape 逻辑、校验约束与执行链路帮助你快速定位源码并理解unpad去 padding 拼接在 Transformer 推理场景中的实际作用。一、Unpad 是什么从 padding 到紧凑拼接在 Transformer 的 decode/推理阶段一个 batch 内不同请求sequence的有效长度往往各不相同。为了按固定形状下发算子常见的做法是统一 padding 到 batch 内最大长度padLength输入张量因此呈现[batch, padLength]的形态其中每个请求真实有效的 token 数记录在seq_len中。Unpad 算子的职责正相反把 padding 过的input_ids重新拆开将各请求的有效 token 依次拼接成一段紧凑的连续序列并在末尾补 0。这一点在参数头文件中被直接写进了结构体注释include/atb/infer_op_params.h对于输入 input_ids把所有有效的 token 拼接在一起并在最后补 0。从知识条目.agent/knowledge/ops/other/unpad/index.md看Unpad 在 ATB 算子体系中属于category: other、tier: S、type: single的单算子执行流水为单阶段由 OpsRunner 直接驱动内核无独立 ACLNN 封装路由文件标注ACLNN: no。二、路由文件速览四文件结构与推荐阅读顺序路由文件.agent/knowledge/routing/unpad.md给出了明确的文件清单与阅读路径这是理解该算子源码的骨架#文件角色1unpad_operation.cppOperation 定义2unpad_operation.hOperation 定义3unpad_ops_runner.cppOps Runner4unpad_ops_runner.hOps Runner推荐按以下顺序阅读每一层关注的重点不同顺序文件重点关注1unpad_operation.h输入输出数量、InferShape 签名2unpad_operation.cppCreateRunner()决策逻辑3unpad_ops_runner.h原生 Ops 执行接口4unpad_ops_runner.cpp原生 Ops 调用链 平台适配对应的源码与配置文件路径均为仓库根目录相对路径Op 目录src/ops/ops_infer/unpad/Kernel 目录src/kernels/mixkernels/unpad参数头文件include/atb/infer_op_params.h三、参数与输入输出语义Unpad 的参数定义极其精简。查看 include/atb/infer_op_params.hinfer::UnpadParam仅包含 8 字节预留字段struct UnpadParam { //! \brief 预留参数 uint8_t rsv[8] {0}; };也就是说 Unpad 是一个无需任何业务参数的纯数据搬运/重组算子所有行为都由输入张量的形状与内容决定。输入张量4 个从 src/ops/ops_infer/unpad/unpad_ops_runner.cpp 的 kernel graph 组装可以看到 4 个输入的命名与角色序号名称形状类型含义0input_ids[batch, padLength]int64padding 后的 token id 序列1cum_offsets_now[batch, 1]int32各请求的累计偏移供内核计算拼接位置2token_num[1, 1]int64本次 batch 的总 token 数3seq_len[batch, 1]int32各请求的有效 token 数输出张量3 个序号名称形状类型含义0x_remove_padding[1, batch × padLength]int64去除 padding 并拼接后的紧凑序列末尾补 01cum_offsets_out[batch, 1]int32输出的累计偏移2padding_offset[1, batch × padLength]int32记录被移除的 padding 的位置偏移四、Operation 层InferShape 与校验约束4.1 类结构与工厂入口src/ops/ops_infer/unpad/unpad_operation.h 定义了UnpadOperation继承自OperationBase重写了GetInputNum()、GetOutputNum()、InferShapeImpl()、CreateRunner()、InferShapeCheckImpl()与SetupCheckImpl()构造时接收infer::UnpadParam。工厂入口在 src/ops/ops_infer/unpad/unpad_operation.cpp通过模板特化CreateOperationinfer::UnpadParam完成参数预留位检查OP_PARAM_RSV_CHECK并new出操作对象template Status CreateOperation(const infer::UnpadParam opParam, Operation **operation) { if (operation nullptr) { return ERROR_INVALID_PARAM; } OP_PARAM_RSV_CHECK(opParam); *operation new (std::nothrow) UnpadOperation(opParam); ... return NO_ERROR; }构造时还会从AtbOperationIrCfg单例中按名称UnpadOperation拉取算子 IR 配置unpad_operation.cpp。4.2 输入输出数量与 InferShape 推导unpad_operation.cpp 通过静态常量固定了张量个数static const uint32_t IN_TENSOR_NUM 4; static const uint32_t OUT_TENSOR_NUM 3;InferShapeImplunpad_operation.cpp的推导规则为三个输出张量的 desc 均以输入 1cum_offsets_now为模板输出 0x_remove_padding与输出 2padding_offset的 dtype 分别被改写为ACL_INT64与保持且形状重置为[1, dim]其中dim是输入 0input_ids所有维度之积——即batch × padLength由于输出形状[1, dim]恰好是输入[batch, padLength]的元素总数这与路由/测试中输出1,300 10×30、1,1600 10×160的形状完全吻合。outTensorDescs.at(0) inTensorDescs.at(1); outTensorDescs.at(1) inTensorDescs.at(1); outTensorDescs.at(DIM_2) inTensorDescs.at(1); outTensorDescs.at(0).dtype ACL_INT64; outTensorDescs.at(0).shape.dims[0] 1; int64_t dim 1; for (size_t i 0; i inTensorDescs.at(0).shape.dimNum; i) { dim * inTensorDescs.at(0).shape.dims[i]; } outTensorDescs.at(0).shape.dims[1] dim; outTensorDescs.at(DIM_2).shape.dims[0] 1; outTensorDescs.at(DIM_2).shape.dims[1] dim;4.3 形状校验InferShapeCheck / SetupCheckInferShapeCheckImpl与SetupCheckImpl的校验逻辑一致unpad_operation.cpp对调用方形成了明确约束这也是使用该算子时必须遵守的潜规则所有 4 个输入张量必须为2 维TensorCheck::IsTensorDescDimNumValid(..., 2)输入 1cum_offsets_now与输入 3seq_len的dim[1]必须为 1即形状应为[batch, 1]输入 2token_num的dim[0]与dim[1]必须均为 1即[1, 1]输入 0input_ids、输入 1cum_offsets_now与输入 3seq_len的dim[0]batch必须一致batch 上限为 64MAX_BATCH_NUM 64超出返回ERROR_INVALID_TENSOR_DIM。4.4 Runner 决策逻辑CreateRunner()unpad_operation.cpp是路由文件中重点标注的决策逻辑——它不经过条件分支而是直接构造UnpadOpsRunnerstd::shared_ptrRunner UnpadOperation::CreateRunner(Context context) const { (void)context; return std::make_sharedUnpadOpsRunner(param_); }五、OpsRunner 层kernel graph 的组装5.1 Runner 类与注册src/ops/ops_infer/unpad/unpad_ops_runner.h 中UnpadOpsRunner继承自OpsRunner仅持有param_成员。其执行接口完全由基类OpsRunner提供因此路由文件将该文件定位为原生 Ops 执行接口。unpad_ops_runner.cpp 末尾通过宏完成两类注册REG_RUNNER_TYPE(UnpadOpsRunner); REG_OP_PARAM(AtbOps::OpParam::Unpad);前者将 Runner 注册进运行时分发机制后者把 ATB 的infer::UnpadParam与内核侧的AtbOps::OpParam::Unpad参数类型关联起来。5.2 调用链组装构造函数unpad_ops_runner.cpp是理解调用链的关键它把 4 个输入、3 个输出一一挂到kernelGraph_上再构造唯一节点节点的opDesc引用内核侧的UnpadOperationAtbOps::OpParam::Unpad unpadParam; unpadNode.opDesc {0, UnpadOperation, unpadParam}; unpadNode.inTensors {inputIds, cumOffsetsNow, tokenNum, seqLen}; unpadNode.outTensors {xRemovePadding, cumOffsetsOut, paddingOffset};这条调用链最终由OpsRunner基类统一走tiling → 内核选择 → AscendC 内核下发的原生 Ops 路径与路由文件标注的 Runner 类型OpsRunner, Operation完全对应。六、Kernel 层tiling 与 AscendC 内核实现6.1 内核侧 Operationsrc/kernels/mixkernels/unpad/unpad_operation.cpp 定义了内核侧同名UnpadOperation继承Mki::OperationBaseGetBestKernel()固定返回名为UnpadKernel的内核GetInputNum()/GetOutputNum()返回 4 入 3 出CheckUnpad()再次校验 batch 一致性、batch ≤ 64InferShapeImpl()与 ATB 层推导规则一致输出 0/2 摊平成[1, Numel(input0)]。6.2 Tiling 参数填充src/kernels/mixkernels/unpad/tiling/unpad_tiling.cpp 负责把形状信息翻译成内核可直接消费的 tiling 数据tilingDataPtr-padLength launchParam.GetInTensor(0).desc.dims[1]; tilingDataPtr-batch launchParam.GetInTensor(0).desc.dims[0]; kernelInfo.SetBlockDim(1);其中padLength取input_ids的dim[1]batch取dim[0]并且固定使用单核blockDim 1执行同时申请 16 字节系统 workspace。6.3 AscendC 内核主流程src/kernels/mixkernels/unpad/op_kernel/unpad.cpp 中的KernelUnpad是实际计算内核几个关键设计编译期常量ELE_PER_BLK 8、ELE_PER_BLK_INT64 4、MAX_BATCH_NUM 64用于将长度按 block 对齐padLengthAlign_、batchAlign_使用TQue队列与TPipe流水缓冲区按BUFFER_NUM 1申请Process()按 batch 逐条处理CopyIn → ComputeOnce计算累计偏移 cumOffsetOut→ ComputeRemovePadding搬运有效段→ CopyOut写出紧凑序列→ CopyOutOnce写出 padding_offset内核核心语义在CopyOutunpad.cpp第 0 个请求直接写到输出开头后续请求写到progress * padLength_ - cumOffsetsBuffer.GetValue(progress - 1)处从而实现跳过 padding、紧凑拼接ClearStep()在末尾把剩余的尾部位置补 0实现参数注释中在最后补 0的行为。该内核通过 src/kernels/mixkernels/unpad/CMakeLists.txt 的add_operation(UnpadOperation ...)注册进内核构建体系。七、测试用例如何验证 Unpad 行为仓库为 Unpad 提供了完整的测试矩阵位于 tests/high_level_test/UnpadOperation/功能用例Boundary_value、Dtype_dataFormat下的UnpadOperation_TestCase.csv以InShape10,30 → OutShape1,300、InShape10,160 → OutShape1,1600验证了输入[batch, padLength]、输出[1, batch×padLength]的摊平语义同时覆盖int64;int32;int64;int32的输入类型组合性能用例Performance/UnpadOperation_Perf_TestCase.csv使用非规则形状如23,33、20,55、52,27等 batch 与 padLength 均不同的组合构造多组 baseline冒烟用例Smoke/UnpadOperation_TestCase.csv以随机形状做泛化验证数据生成范围为-100,100。此外Unpad 的拼接思路在flash_attention、fa_mla等融合算子的测试中也被大量复用见 tests/apitest/kernelstest/mix/test_flash_attention.py 中# unpad encoder注释印证了unpad 是动态 batch 注意力计算前置整理步骤的定位。八、小结一条从路由到内核的完整链路回到路由文件Unpad 算子的源码阅读路径可以总结为一条四层链路入口src/ops/ops_infer/unpad/unpad_operation.cpp — 工厂创建、InferShape、形状校验、CreateRunner()决策组装src/ops/ops_infer/unpad/unpad_ops_runner.cpp — 将 4 入 3 出挂接为 kernel graph 节点并注册 Runner内核侧src/kernels/mixkernels/unpad/unpad_operation.cpp — 选择UnpadKernel计算src/kernels/mixkernels/unpad/op_kernel/unpad.cpp 与 tiling/unpad_tiling.cpp — 单核 tiling AscendC 内核完成去 padding 拼接与补 0。掌握了这条链路你在阅读flash_attention、paged_attention等更复杂的动态 batch 算子时也能快速识别出其中复用 Unpad 思路的预处理/后处理环节。相关路由索引可继续参考 .agent/knowledge/README.md 与 .agent/knowledge/ops/other/unpad/index.md。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表