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）给出了明确的文件清单与阅读路径，这是理解该算子源码的骨架：
| # | 文件 | 角色 |
|---|---|---|
| 1 | unpad_operation.cpp | Operation 定义 |
| 2 | unpad_operation.h | Operation 定义 |
| 3 | unpad_ops_runner.cpp | Ops Runner |
| 4 | unpad_ops_runner.h | Ops Runner |
推荐按以下顺序阅读，每一层关注的重点不同：
| 顺序 | 文件 | 重点关注 |
|---|---|---|
| 1 | unpad_operation.h | 输入输出数量、InferShape 签名 |
| 2 | unpad_operation.cpp | CreateRunner()决策逻辑 |
| 3 | unpad_ops_runner.h | 原生 Ops 执行接口 |
| 4 | unpad_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.h，infer::UnpadParam仅包含 8 字节预留字段：
struct UnpadParam { //! \brief 预留参数 uint8_t rsv[8] = {0}; };也就是说 Unpad 是一个无需任何业务参数的纯数据搬运/重组算子，所有行为都由输入张量的形状与内容决定。
输入张量（4 个）
从 src/ops/ops_infer/unpad/unpad_ops_runner.cpp 的 kernel graph 组装可以看到 4 个输入的命名与角色：
| 序号 | 名称 | 形状 | 类型 | 含义 |
|---|---|---|---|---|
| 0 | input_ids | [batch, padLength] | int64 | padding 后的 token id 序列 |
| 1 | cum_offsets_now | [batch, 1] | int32 | 各请求的累计偏移（供内核计算拼接位置） |
| 2 | token_num | [1, 1] | int64 | 本次 batch 的总 token 数 |
| 3 | seq_len | [batch, 1] | int32 | 各请求的有效 token 数 |
输出张量（3 个）
| 序号 | 名称 | 形状 | 类型 | 含义 |
|---|---|---|---|---|
| 0 | x_remove_padding | [1, batch × padLength] | int64 | 去除 padding 并拼接后的紧凑序列（末尾补 0） |
| 1 | cum_offsets_out | [batch, 1] | int32 | 输出的累计偏移 |
| 2 | padding_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，通过模板特化CreateOperation<infer::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;InferShapeImpl（unpad_operation.cpp）的推导规则为：
- 三个输出张量的 desc 均以输入 1（
cum_offsets_now）为模板； - 输出 0（
x_remove_padding）与输出 2（padding_offset）的 dtype 分别被改写为ACL_INT64与保持，且形状重置为[1, dim]，其中dim是输入 0（input_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 / SetupCheck）
InferShapeCheckImpl与SetupCheckImpl的校验逻辑一致（unpad_operation.cpp），对调用方形成了明确约束，这也是使用该算子时必须遵守的"潜规则"：
- 所有 4 个输入张量必须为2 维（
TensorCheck::IsTensorDescDimNumValid(..., 2)）； - 输入 1（
cum_offsets_now）与输入 3（seq_len）的dim[1]必须为 1，即形状应为[batch, 1]； - 输入 2（
token_num）的dim[0]与dim[1]必须均为 1，即[1, 1]； - 输入 0（
input_ids）、输入 1（cum_offsets_now）与输入 3（seq_len）的dim[0]（batch）必须一致； - batch 上限为 64（
MAX_BATCH_NUM = 64），超出返回ERROR_INVALID_TENSOR_DIM。
4.4 Runner 决策逻辑
CreateRunner()（unpad_operation.cpp）是路由文件中重点标注的"决策逻辑"——它不经过条件分支，而是直接构造UnpadOpsRunner：
std::shared_ptr<Runner> UnpadOperation::CreateRunner(Context &context) const { (void)context; return std::make_shared<UnpadOpsRunner>(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引用内核侧的"UnpadOperation"：
AtbOps::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 内核侧 Operation
src/kernels/mixkernels/unpad/unpad_operation.cpp 定义了内核侧同名UnpadOperation，继承Mki::OperationBase：
GetBestKernel()固定返回名为"UnpadKernel"的内核；GetInputNum()/GetOutputNum()返回 4 入 3 出；CheckUnpad()再次校验 batch 一致性、batch ≤ 64；InferShapeImpl()与 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）；- 内核核心语义在
CopyOut（unpad.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）：以InShape=10,30 → OutShape=1,300、InShape=10,160 → OutShape=1,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），仅供参考
