尧图网站建设 
ascend-transformer-boost 训练算子 unpad_with_hidden_state 源码全解析去填充隐藏状态算子的实现与调用链【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost导读unpad_with_hidden_state是 ascend-transformer-boostCANN 基于昇腾 AI 处理器的 Transformer 加速库中面向训练侧的去填充unpad融合算子其职责是把按maxSeqLen补齐padding后的批式隐藏状态序列按每条样本的真实长度压缩回紧凑的连续布局从而为后续按真实序列长度计算如反向传播、loss 计算提供紧凑输入。本文以仓库内路由文档为主线逐层拆解该算子的文件组织、参数结构、Operation 层 InferShape、OpsRunner 图构建以及 Kernel 层的多核分片实现读完即可掌握该算子的完整调用链与源码阅读方法。1. 算子路由定位先读路由文件再读知识条目在.agent/knowledge/目录下仓库维护了一套面向 Agent 的算子知识索引。定位任意算子的标准路径是1. 在 .agent/knowledge/README.md 搜索 Op 名 → 获取 tier 与 category 2. 读 .agent/knowledge/routing/{op_name}.md → 获取文件清单和阅读顺序 3. 按需读 .agent/knowledge/ops/{category}/{op_name}/index.md → 获取完整知识unpad_with_hidden_state在路由文件中登记的信息如下属性值分类categorytrain训练算子复杂度tierS简单算子4 个文件单一路径Runner 类型OpsRunner, OperationACLNN 支持no纯原生 Ops 路径无 ACLNN 封装预估阅读时间3-5 分钟在知识库主索引中该算子被归入train 分类训练算子共 10 个与纯 Ops 路径57 个不经过 ACLNN、直接走原生 Ops API 的算子之一。其详细知识条目位于.agent/knowledge/ops/train/unpad_with_hidden_state/index.md该条目当前状态为complete登记了op.name unpad_with_hidden_state、category train、tier S、type single以及源码路径src/ops/ops_train/unpad_with_hidden_state/。路由文件给出的推荐阅读顺序本身就是一条完整的源码学习路径顺序文件重点关注1unpad_with_hidden_state_operation.h输入输出数量、InferShape 签名2unpad_with_hidden_state_operation.cppCreateRunner()决策逻辑3unpad_with_hidden_state_ops_runner.h原生 Ops 执行接口4unpad_with_hidden_state_ops_runner.cpp原生 Ops 调用链 平台适配2. 文件清单与源码路径路由文件列出该算子共4 个文件全部位于 src/ops/ops_train/unpad_with_hidden_state/#文件角色1unpad_with_hidden_state_operation.cppOperation 定义2unpad_with_hidden_state_operation.hOperation 定义3unpad_with_hidden_state_ops_runner.cppOps Runner4unpad_with_hidden_state_ops_runner.hOps Runner对应的 Kernel 目录为 src/kernels/mixkernels/unpad_with_hidden_state其中包含unpad_with_hidden_state_kernel.cppKernel 注册入口REG_KERNEL_BASE(UnpadWithHiddenStateKernel)tiling/unpad_with_hidden_state_tiling.cpp/.hTiling 计算多核任务切分tiling/tiling_data.hTiling 数据结构UnpadWithHiddenStateTilingDataop_kernel/unpad_with_hidden_state.cppAscendC 算子内核实现CopyIn / Compute / CopyOut参数头文件为 include/atb/train_op_params.h——路由文件指出训练算子train的参数统一声明在train_op_params.h而推理算子infer的参数则位于include/atb/infer_op_params.h这是定位参数结构体时的约定。从目录结构看src/ops/ops_train/下与unpad_with_hidden_state并列的还有pad_with_hidden_state、rms_norm_backward、rope_grad、strided_batch_matmul等训练算子二者构成一对先补齐、后压缩的配套操作。3. 参数结构UnpadWithHiddenStateParam算子的参数结构体定义在 include/atb/train_op_params.h 的UnpadWithHiddenStateParam约 L196-L225关键字段如下//! \struct UnpadWithHiddenStateParam struct UnpadWithHiddenStateParam { //! \brief 每条样本的真实序列长度batch 个元素的数组元素为正整数且不超过 maxSeqLen。 std::vectorint32_t qSeqLen; //! \brief 最大 qSeqLen 值。取值不超过 4096。 int32_t maxSeqLen 4096; }; inline bool operator(const UnpadWithHiddenStateParam left, const UnpadWithHiddenStateParam right) { return left.qSeqLen right.qSeqLen left.maxSeqLen right.maxSeqLen; }参数语义qSeqLen必需长度等于 batch size 的整型数组表示批内每条样本的真实序列长度。它的作用是告知算子每条样本实际需要保留多少个 token从而精确计算输入偏移与输出偏移。maxSeqLen可选默认 4096批内统一补齐的目标序列长度即输入张量第 1 维的大小。校验要求0 maxSeqLen 4096。结构体实现了operator供 Runner 的SetParam做参数变更检测详见第 5 节。4. Operation 层InferShape 与参数校验4.1 类声明与输入输出数量unpad_with_hidden_state_operation.h 中UnpadWithHiddenStateOperation继承自OperationBase暴露以下接口class UnpadWithHiddenStateOperation : public OperationBase { public: explicit UnpadWithHiddenStateOperation(const train::UnpadWithHiddenStateParam param); uint32_t GetInputNum() const override; uint32_t GetOutputNum() const override; train::UnpadWithHiddenStateParam GetParam() const; void SetParam(const train::UnpadWithHiddenStateParam param); protected: Status InferShapeImpl(const SVectorTensorDesc inTensorDescs, SVectorTensorDesc outTensorDescs) const override; std::shared_ptrRunner CreateRunner(Context context) const override; Status InferShapeCheckImpl(const SVectorTensorDesc inTensorDescs) const override; Status SetupCheckImpl(const SVectorTensor inTensors, const SVectorTensor outTensors) const override; private: train::UnpadWithHiddenStateParam param_; };在 unpad_with_hidden_state_operation.cpp 中输入输出数量固定为1 输入 1 输出GetInputNum()返回 1GetOutputNum()返回 1。4.2 平台与参数校验ParamCheckoperation.cpp 中的ParamCheck是进入执行前的第一道关卡bool ParamCheck(const atb::train::UnpadWithHiddenStateParam opParam) { if (!atb::GetSingletonatb::Config().Is910B()) { ATB_LOG(ERROR) UnpadWithHiddenStateOperation is only supported in Atlas 800I A2 inference product.; return false; } if (opParam.maxSeqLen 0 || opParam.maxSeqLen MAX_SEQLEN) { ATB_LOG(ERROR) UnpadWithHiddenStateParam maxSeqLen should be in (0, MAX_SEQLEN ], but get opParam.maxSeqLen; return false; } return atb::OperationUtil::QSeqLenCheck(opParam.qSeqLen, opParam.maxSeqLen); }三个校验要点平台限制仅支持 Atlas 800I A2910B推理产品形态其他平台直接报错——这是源码中明确写死的平台约束MAX_SEQLEN 4096。maxSeqLen 范围必须满足0 maxSeqLen 4096。qSeqLen 合法性通过OperationUtil::QSeqLenCheck校验逐元素确保每条样本的真实长度落在(0, maxSeqLen]区间。4.3 张量形状约束输入输出张量的形状校验同样在 operation.cpp 中输入张量CheckInTensorDescsL38-L53要求 3 维inDimNum 3且满足dims[0] qSeqLen.size()batch 大小与样本数一致dims[1] maxSeqLen补齐后的序列长度。输出张量CheckOutTensorsL62-L81要求 2 维outDimNum 2且满足dims[0] sum(qSeqLen)各样本真实长度之和dims[1] hiddenSize与输入第 2 维一致。4.4 InferShapeImpl 的输出推导InferShapeImpl 在形状校验通过后计算输出描述int64_t outerSize 0; for (auto sampleSeqLen : param_.qSeqLen) { outerSize sampleSeqLen; // 输出第 0 维 Σ qSeqLen } int64_t hiddenSize inTensorDescs.at(0).shape.dims[DIM_2]; outDesc.dtype inTensorDescs.at(0).dtype; // 数据类型、格式与输入保持一致 outDesc.format inTensorDescs.at(0).format; outDesc.shape.dimNum outDimNum; // 2 维 outDesc.shape.dims[0] outerSize; outDesc.shape.dims[1] hiddenSize;即输入[batch, maxSeqLen, hiddenSize]→ 输出[ΣqSeqLen, hiddenSize]dtype 与 format 原样透传。这正是unpad去填充的语义把 padding 占位去掉只保留真实 token同时保持 hidden 维不变。5. OpsRunner 层原生 Ops 图构建与调用链5.1 Runner 类型注册与选择在 CreateRunner 中Operation 通过RunnerTypeRegister::GetRunnerTypeIdx(UnpadWithHiddenStateOpsRunner)获取 Runner 类型索引再从ContextBase的RunnerPool中复用或新建UnpadWithHiddenStateOpsRunner。这就是路由文件中 Runner 类型: OpsRunner, Operation 的含义——该算子没有 ACLNN 路径执行统一走 OpsRunner。Runner 自身在 unpad_with_hidden_state_ops_runner.cpp 末尾完成注册REG_RUNNER_TYPE(UnpadWithHiddenStateOpsRunner); REG_OP_PARAM(AtbOps::OpParam::UnpadWithHiddenState);5.2 SetupKernelGraph构建 Mki Kernel GraphSetupKernelGraph 是 OpsRunner 的核心逻辑负责把 ATB 层的参数转换为 Mki Kernel 图节点Status UnpadWithHiddenStateOpsRunner::SetupKernelGraph(const OpsTensorPack opsTensorPack) { const size_t inTensorSize 1; const size_t outTensorSize 1; kernelGraph_.inTensors.resize(inTensorSize); // 1 个输入 tensor kernelGraph_.outTensors.resize(outTensorSize); // 1 个输出 tensor Mki::Tensor dataInputTensor kernelGraph_.inTensors.at(0); Mki::Tensor dataOutputTensor kernelGraph_.outTensors.at(0); kernelGraph_.nodes.resize(1); auto trainUnpadNode kernelGraph_.nodes.at(0); AtbOps::OpParam::UnpadWithHiddenState unpadWithHiddenStateParam; unpadWithHiddenStateParam.qSeqLen param_.qSeqLen; unpadWithHiddenStateParam.maxSeqLen static_castuint32_t(param_.maxSeqLen); trainUnpadNode.opDesc {0, UnpadWithHiddenStateOperation, unpadWithHiddenStateParam}; trainUnpadNode.inTensors {dataInputTensor}; trainUnpadNode.outTensors {dataOutputTensor}; return NO_ERROR; }该函数完成三件事把 ATB 层train::UnpadWithHiddenStateParam映射为 Kernel 层AtbOps::OpParam::UnpadWithHiddenStateqSeqLen、maxSeqLen一一对应注册一个名为UnpadWithHiddenStateOperation的图节点把唯一的输入输出 tensor 挂接到节点上。图构建完成后由基类 src/atb/runner/ops_runner.h 负责启动执行。5.3 SetParam参数热更新SetParam 利用第 3 节提到的operator检测参数是否变化仅在变化时更新内部param_并置isParamUpdated_ true从而触发 Kernel 图重建避免无效重建开销void UnpadWithHiddenStateOpsRunner::SetParam(const Mki::Any param) { train::UnpadWithHiddenStateParam newParam Mki::AnyCasttrain::UnpadWithHiddenStateParam(param); if (!(newParam param_)) { ATB_LOG(DEBUG) UnpadWithHiddenStateOpsRunner Param Changed!; param_ newParam; isParamUpdated_ true; } }6. Kernel 层Tiling 多核分片与 AscendC 内核6.1 Kernel 注册入口unpad_with_hidden_state_kernel.cpp 定义了UnpadWithHiddenStateKernel : KernelBase其中CanSupport校验 1 输入、1 输出且参数类型为OpParam::UnpadWithHiddenStateGetTilingSize返回sizeof(UnpadWithHiddenStateTilingData)InitImpl调用UnpadWithHiddenStateTiling(launchParam, kernelInfo_)生成 Tiling 数据。6.2 Tiling按样本切分 多核负载均衡Tiling 实现L39-L90的核心逻辑获取向量核数量coreNum、batch 大小、maxSeqLen 与 hiddenSize输入第 2 维校验batchSize MAX_BATCH_SIZE、maxSeqLen SEQLEN_LIMIT计算单核单次搬运的数据块bufferSize AlignDownToBasicBlock(UB_SIZE / BUFFER_NUM / UNPAD_WITH_HIDDEN_STATE_COMPUTE_DIM)tileLength bufferSize / ELEMENT_SIZE逐样本计算偏移对第sampleIndex条样本输入偏移inputOffset maxSeqLen * hiddenSize按补齐后长度步进输出偏移outputOffset sampleSeqLen * hiddenSize按真实长度步进并把inputOffset / outputOffset / tileInfo写入 TilingData对每条样本的数据量sampleDataLength sampleSeqLen * hiddenSize调用GetTileInfo做多核均分uint32_t tileNum dataLength / tileLength; tileInfo.formerCoreNum tileNum % coreNum; // 多承担 1 块的前缀核数 tileInfo.formerCoreTileNum tileNum / coreNum 1; // 前缀核每核承担的块数 tileInfo.lastTileLength dataLength % tileLength; // 末尾不足一块的剩余长度6.3 AscendC 内核逐样本 CopyIn → Compute → CopyOut算子内核使用 AscendC 编写代码中以half类型为例InitpadDataGm指向补齐后的输入全局内存unpadDataGm指向输出并初始化双缓冲队列BUFFER_NUM个 bufferProcess外层按样本循环每轮从 TilingData 读取该样本的inputOffset / outputOffset / tileInfo内层依据核号GetBlockIdx()计算本核承担的首块偏移innerOffset与块数tileNum逐块执行CopyIn → Compute → CopyOut尾块处理lastTileLength 0的剩余数据仅由最后一个核isLastCore负责搬运CopyIn 的边界处理当搬运长度不是ELEMENT_PER_BASIC_BLOCK的整数倍时使用DataCopyPad进行带 padding 的搬运保证非对齐场景下数据正确落位。整体上Kernel 层的本质是在 GM 到 UB 之间做一次跳过 padding 的连续搬运从第sampleIndex个样本的maxSeqLen * hiddenSize步进区域中只抽取前qSeqLen[sampleIndex] * hiddenSize个元素按紧凑偏移写入输出。7. 总结一次 unpad 的完整调用链从路由文档与源码可归纳出unpad_with_hidden_state的完整执行链路ATB API构造 UnpadWithHiddenStateOperation → InferShapeCheckImpl / SetupCheckImpl平台 参数 张量形状校验 → InferShapeImpl[batch, maxSeqLen, hidden] → [ΣqSeqLen, hidden] → CreateRunnerRunnerPool 复用 UnpadWithHiddenStateOpsRunner → SetupKernelGraphATB 参数 → Mki Kernel 图节点 → UnpadWithHiddenStateTiling逐样本偏移计算 多核 tile 切分 → AscendC Kernel多核并行 CopyIn → Compute → CopyOut跳过 padding → 输出紧凑的 [ΣqSeqLen, hiddenSize] 张量该算子的源码阅读入口汇总路由与知识索引.agent/knowledge/routing/unpad_with_hidden_state.md、.agent/knowledge/ops/train/unpad_with_hidden_state/index.md、.agent/knowledge/README.mdOperation 层unpad_with_hidden_state_operation.h、unpad_with_hidden_state_operation.cppRunner 层unpad_with_hidden_state_ops_runner.h、unpad_with_hidden_state_ops_runner.cppKernel 层unpad_with_hidden_state_kernel.cpp、tiling/unpad_with_hidden_state_tiling.cpp、op_kernel/unpad_with_hidden_state.cpp参数定义include/atb/train_op_params.h对于希望向 ascend-transformer-boost 贡献或调试训练算子的开发者建议严格遵循路由文件给出的四步阅读顺序先看operation.h明确接口签名再看operation.cpp理解校验与 Runner 决策最后进入ops_runner与 Kernel 层追踪实际执行即可在 3-5 分钟内建立对该算子的完整认知。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ascend-transformer-boost 训练算子 unpad_with_hidden_state 源码全解析：去填充隐藏状态算子的实现与调用链
ascend-transformer-boost 训练算子 unpad_with_hidden_state 源码全解析去填充隐藏状态算子的实现与调用链【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost导读unpad_with_hidden_state是 ascend-transformer-boostCANN 基于昇腾 AI 处理器的 Transformer 加速库中面向训练侧的去填充unpad融合算子其职责是把按maxSeqLen补齐padding后的批式隐藏状态序列按每条样本的真实长度压缩回紧凑的连续布局从而为后续按真实序列长度计算如反向传播、loss 计算提供紧凑输入。本文以仓库内路由文档为主线逐层拆解该算子的文件组织、参数结构、Operation 层 InferShape、OpsRunner 图构建以及 Kernel 层的多核分片实现读完即可掌握该算子的完整调用链与源码阅读方法。1. 算子路由定位先读路由文件再读知识条目在.agent/knowledge/目录下仓库维护了一套面向 Agent 的算子知识索引。定位任意算子的标准路径是1. 在 .agent/knowledge/README.md 搜索 Op 名 → 获取 tier 与 category 2. 读 .agent/knowledge/routing/{op_name}.md → 获取文件清单和阅读顺序 3. 按需读 .agent/knowledge/ops/{category}/{op_name}/index.md → 获取完整知识unpad_with_hidden_state在路由文件中登记的信息如下属性值分类categorytrain训练算子复杂度tierS简单算子4 个文件单一路径Runner 类型OpsRunner, OperationACLNN 支持no纯原生 Ops 路径无 ACLNN 封装预估阅读时间3-5 分钟在知识库主索引中该算子被归入train 分类训练算子共 10 个与纯 Ops 路径57 个不经过 ACLNN、直接走原生 Ops API 的算子之一。其详细知识条目位于.agent/knowledge/ops/train/unpad_with_hidden_state/index.md该条目当前状态为complete登记了op.name unpad_with_hidden_state、category train、tier S、type single以及源码路径src/ops/ops_train/unpad_with_hidden_state/。路由文件给出的推荐阅读顺序本身就是一条完整的源码学习路径顺序文件重点关注1unpad_with_hidden_state_operation.h输入输出数量、InferShape 签名2unpad_with_hidden_state_operation.cppCreateRunner()决策逻辑3unpad_with_hidden_state_ops_runner.h原生 Ops 执行接口4unpad_with_hidden_state_ops_runner.cpp原生 Ops 调用链 平台适配2. 文件清单与源码路径路由文件列出该算子共4 个文件全部位于 src/ops/ops_train/unpad_with_hidden_state/#文件角色1unpad_with_hidden_state_operation.cppOperation 定义2unpad_with_hidden_state_operation.hOperation 定义3unpad_with_hidden_state_ops_runner.cppOps Runner4unpad_with_hidden_state_ops_runner.hOps Runner对应的 Kernel 目录为 src/kernels/mixkernels/unpad_with_hidden_state其中包含unpad_with_hidden_state_kernel.cppKernel 注册入口REG_KERNEL_BASE(UnpadWithHiddenStateKernel)tiling/unpad_with_hidden_state_tiling.cpp/.hTiling 计算多核任务切分tiling/tiling_data.hTiling 数据结构UnpadWithHiddenStateTilingDataop_kernel/unpad_with_hidden_state.cppAscendC 算子内核实现CopyIn / Compute / CopyOut参数头文件为 include/atb/train_op_params.h——路由文件指出训练算子train的参数统一声明在train_op_params.h而推理算子infer的参数则位于include/atb/infer_op_params.h这是定位参数结构体时的约定。从目录结构看src/ops/ops_train/下与unpad_with_hidden_state并列的还有pad_with_hidden_state、rms_norm_backward、rope_grad、strided_batch_matmul等训练算子二者构成一对先补齐、后压缩的配套操作。3. 参数结构UnpadWithHiddenStateParam算子的参数结构体定义在 include/atb/train_op_params.h 的UnpadWithHiddenStateParam约 L196-L225关键字段如下//! \struct UnpadWithHiddenStateParam struct UnpadWithHiddenStateParam { //! \brief 每条样本的真实序列长度batch 个元素的数组元素为正整数且不超过 maxSeqLen。 std::vectorint32_t qSeqLen; //! \brief 最大 qSeqLen 值。取值不超过 4096。 int32_t maxSeqLen 4096; }; inline bool operator(const UnpadWithHiddenStateParam left, const UnpadWithHiddenStateParam right) { return left.qSeqLen right.qSeqLen left.maxSeqLen right.maxSeqLen; }参数语义qSeqLen必需长度等于 batch size 的整型数组表示批内每条样本的真实序列长度。它的作用是告知算子每条样本实际需要保留多少个 token从而精确计算输入偏移与输出偏移。maxSeqLen可选默认 4096批内统一补齐的目标序列长度即输入张量第 1 维的大小。校验要求0 maxSeqLen 4096。结构体实现了operator供 Runner 的SetParam做参数变更检测详见第 5 节。4. Operation 层InferShape 与参数校验4.1 类声明与输入输出数量unpad_with_hidden_state_operation.h 中UnpadWithHiddenStateOperation继承自OperationBase暴露以下接口class UnpadWithHiddenStateOperation : public OperationBase { public: explicit UnpadWithHiddenStateOperation(const train::UnpadWithHiddenStateParam param); uint32_t GetInputNum() const override; uint32_t GetOutputNum() const override; train::UnpadWithHiddenStateParam GetParam() const; void SetParam(const train::UnpadWithHiddenStateParam param); protected: Status InferShapeImpl(const SVectorTensorDesc inTensorDescs, SVectorTensorDesc outTensorDescs) const override; std::shared_ptrRunner CreateRunner(Context context) const override; Status InferShapeCheckImpl(const SVectorTensorDesc inTensorDescs) const override; Status SetupCheckImpl(const SVectorTensor inTensors, const SVectorTensor outTensors) const override; private: train::UnpadWithHiddenStateParam param_; };在 unpad_with_hidden_state_operation.cpp 中输入输出数量固定为1 输入 1 输出GetInputNum()返回 1GetOutputNum()返回 1。4.2 平台与参数校验ParamCheckoperation.cpp 中的ParamCheck是进入执行前的第一道关卡bool ParamCheck(const atb::train::UnpadWithHiddenStateParam opParam) { if (!atb::GetSingletonatb::Config().Is910B()) { ATB_LOG(ERROR) UnpadWithHiddenStateOperation is only supported in Atlas 800I A2 inference product.; return false; } if (opParam.maxSeqLen 0 || opParam.maxSeqLen MAX_SEQLEN) { ATB_LOG(ERROR) UnpadWithHiddenStateParam maxSeqLen should be in (0, MAX_SEQLEN ], but get opParam.maxSeqLen; return false; } return atb::OperationUtil::QSeqLenCheck(opParam.qSeqLen, opParam.maxSeqLen); }三个校验要点平台限制仅支持 Atlas 800I A2910B推理产品形态其他平台直接报错——这是源码中明确写死的平台约束MAX_SEQLEN 4096。maxSeqLen 范围必须满足0 maxSeqLen 4096。qSeqLen 合法性通过OperationUtil::QSeqLenCheck校验逐元素确保每条样本的真实长度落在(0, maxSeqLen]区间。4.3 张量形状约束输入输出张量的形状校验同样在 operation.cpp 中输入张量CheckInTensorDescsL38-L53要求 3 维inDimNum 3且满足dims[0] qSeqLen.size()batch 大小与样本数一致dims[1] maxSeqLen补齐后的序列长度。输出张量CheckOutTensorsL62-L81要求 2 维outDimNum 2且满足dims[0] sum(qSeqLen)各样本真实长度之和dims[1] hiddenSize与输入第 2 维一致。4.4 InferShapeImpl 的输出推导InferShapeImpl 在形状校验通过后计算输出描述int64_t outerSize 0; for (auto sampleSeqLen : param_.qSeqLen) { outerSize sampleSeqLen; // 输出第 0 维 Σ qSeqLen } int64_t hiddenSize inTensorDescs.at(0).shape.dims[DIM_2]; outDesc.dtype inTensorDescs.at(0).dtype; // 数据类型、格式与输入保持一致 outDesc.format inTensorDescs.at(0).format; outDesc.shape.dimNum outDimNum; // 2 维 outDesc.shape.dims[0] outerSize; outDesc.shape.dims[1] hiddenSize;即输入[batch, maxSeqLen, hiddenSize]→ 输出[ΣqSeqLen, hiddenSize]dtype 与 format 原样透传。这正是unpad去填充的语义把 padding 占位去掉只保留真实 token同时保持 hidden 维不变。5. OpsRunner 层原生 Ops 图构建与调用链5.1 Runner 类型注册与选择在 CreateRunner 中Operation 通过RunnerTypeRegister::GetRunnerTypeIdx(UnpadWithHiddenStateOpsRunner)获取 Runner 类型索引再从ContextBase的RunnerPool中复用或新建UnpadWithHiddenStateOpsRunner。这就是路由文件中 Runner 类型: OpsRunner, Operation 的含义——该算子没有 ACLNN 路径执行统一走 OpsRunner。Runner 自身在 unpad_with_hidden_state_ops_runner.cpp 末尾完成注册REG_RUNNER_TYPE(UnpadWithHiddenStateOpsRunner); REG_OP_PARAM(AtbOps::OpParam::UnpadWithHiddenState);5.2 SetupKernelGraph构建 Mki Kernel GraphSetupKernelGraph 是 OpsRunner 的核心逻辑负责把 ATB 层的参数转换为 Mki Kernel 图节点Status UnpadWithHiddenStateOpsRunner::SetupKernelGraph(const OpsTensorPack opsTensorPack) { const size_t inTensorSize 1; const size_t outTensorSize 1; kernelGraph_.inTensors.resize(inTensorSize); // 1 个输入 tensor kernelGraph_.outTensors.resize(outTensorSize); // 1 个输出 tensor Mki::Tensor dataInputTensor kernelGraph_.inTensors.at(0); Mki::Tensor dataOutputTensor kernelGraph_.outTensors.at(0); kernelGraph_.nodes.resize(1); auto trainUnpadNode kernelGraph_.nodes.at(0); AtbOps::OpParam::UnpadWithHiddenState unpadWithHiddenStateParam; unpadWithHiddenStateParam.qSeqLen param_.qSeqLen; unpadWithHiddenStateParam.maxSeqLen static_castuint32_t(param_.maxSeqLen); trainUnpadNode.opDesc {0, UnpadWithHiddenStateOperation, unpadWithHiddenStateParam}; trainUnpadNode.inTensors {dataInputTensor}; trainUnpadNode.outTensors {dataOutputTensor}; return NO_ERROR; }该函数完成三件事把 ATB 层train::UnpadWithHiddenStateParam映射为 Kernel 层AtbOps::OpParam::UnpadWithHiddenStateqSeqLen、maxSeqLen一一对应注册一个名为UnpadWithHiddenStateOperation的图节点把唯一的输入输出 tensor 挂接到节点上。图构建完成后由基类 src/atb/runner/ops_runner.h 负责启动执行。5.3 SetParam参数热更新SetParam 利用第 3 节提到的operator检测参数是否变化仅在变化时更新内部param_并置isParamUpdated_ true从而触发 Kernel 图重建避免无效重建开销void UnpadWithHiddenStateOpsRunner::SetParam(const Mki::Any param) { train::UnpadWithHiddenStateParam newParam Mki::AnyCasttrain::UnpadWithHiddenStateParam(param); if (!(newParam param_)) { ATB_LOG(DEBUG) UnpadWithHiddenStateOpsRunner Param Changed!; param_ newParam; isParamUpdated_ true; } }6. Kernel 层Tiling 多核分片与 AscendC 内核6.1 Kernel 注册入口unpad_with_hidden_state_kernel.cpp 定义了UnpadWithHiddenStateKernel : KernelBase其中CanSupport校验 1 输入、1 输出且参数类型为OpParam::UnpadWithHiddenStateGetTilingSize返回sizeof(UnpadWithHiddenStateTilingData)InitImpl调用UnpadWithHiddenStateTiling(launchParam, kernelInfo_)生成 Tiling 数据。6.2 Tiling按样本切分 多核负载均衡Tiling 实现L39-L90的核心逻辑获取向量核数量coreNum、batch 大小、maxSeqLen 与 hiddenSize输入第 2 维校验batchSize MAX_BATCH_SIZE、maxSeqLen SEQLEN_LIMIT计算单核单次搬运的数据块bufferSize AlignDownToBasicBlock(UB_SIZE / BUFFER_NUM / UNPAD_WITH_HIDDEN_STATE_COMPUTE_DIM)tileLength bufferSize / ELEMENT_SIZE逐样本计算偏移对第sampleIndex条样本输入偏移inputOffset maxSeqLen * hiddenSize按补齐后长度步进输出偏移outputOffset sampleSeqLen * hiddenSize按真实长度步进并把inputOffset / outputOffset / tileInfo写入 TilingData对每条样本的数据量sampleDataLength sampleSeqLen * hiddenSize调用GetTileInfo做多核均分uint32_t tileNum dataLength / tileLength; tileInfo.formerCoreNum tileNum % coreNum; // 多承担 1 块的前缀核数 tileInfo.formerCoreTileNum tileNum / coreNum 1; // 前缀核每核承担的块数 tileInfo.lastTileLength dataLength % tileLength; // 末尾不足一块的剩余长度6.3 AscendC 内核逐样本 CopyIn → Compute → CopyOut算子内核使用 AscendC 编写代码中以half类型为例InitpadDataGm指向补齐后的输入全局内存unpadDataGm指向输出并初始化双缓冲队列BUFFER_NUM个 bufferProcess外层按样本循环每轮从 TilingData 读取该样本的inputOffset / outputOffset / tileInfo内层依据核号GetBlockIdx()计算本核承担的首块偏移innerOffset与块数tileNum逐块执行CopyIn → Compute → CopyOut尾块处理lastTileLength 0的剩余数据仅由最后一个核isLastCore负责搬运CopyIn 的边界处理当搬运长度不是ELEMENT_PER_BASIC_BLOCK的整数倍时使用DataCopyPad进行带 padding 的搬运保证非对齐场景下数据正确落位。整体上Kernel 层的本质是在 GM 到 UB 之间做一次跳过 padding 的连续搬运从第sampleIndex个样本的maxSeqLen * hiddenSize步进区域中只抽取前qSeqLen[sampleIndex] * hiddenSize个元素按紧凑偏移写入输出。7. 总结一次 unpad 的完整调用链从路由文档与源码可归纳出unpad_with_hidden_state的完整执行链路ATB API构造 UnpadWithHiddenStateOperation → InferShapeCheckImpl / SetupCheckImpl平台 参数 张量形状校验 → InferShapeImpl[batch, maxSeqLen, hidden] → [ΣqSeqLen, hidden] → CreateRunnerRunnerPool 复用 UnpadWithHiddenStateOpsRunner → SetupKernelGraphATB 参数 → Mki Kernel 图节点 → UnpadWithHiddenStateTiling逐样本偏移计算 多核 tile 切分 → AscendC Kernel多核并行 CopyIn → Compute → CopyOut跳过 padding → 输出紧凑的 [ΣqSeqLen, hiddenSize] 张量该算子的源码阅读入口汇总路由与知识索引.agent/knowledge/routing/unpad_with_hidden_state.md、.agent/knowledge/ops/train/unpad_with_hidden_state/index.md、.agent/knowledge/README.mdOperation 层unpad_with_hidden_state_operation.h、unpad_with_hidden_state_operation.cppRunner 层unpad_with_hidden_state_ops_runner.h、unpad_with_hidden_state_ops_runner.cppKernel 层unpad_with_hidden_state_kernel.cpp、tiling/unpad_with_hidden_state_tiling.cpp、op_kernel/unpad_with_hidden_state.cpp参数定义include/atb/train_op_params.h对于希望向 ascend-transformer-boost 贡献或调试训练算子的开发者建议严格遵循路由文件给出的四步阅读顺序先看operation.h明确接口签名再看operation.cpp理解校验与 Runner 决策最后进入ops_runner与 Kernel 层追踪实际执行即可在 3-5 分钟内建立对该算子的完整认知。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考