ascend-transformer-boost 的 KvCacheOperation 源码导读从路由文件到 Ops Runner 调用链【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本文以仓库中 kv_cache 路由文件 为主线深入解析 ascend-transformer-boostATB中推理场景的 KV Cache 写入算子KvCacheOperation。读者将掌握该算子的文件组织、推荐阅读顺序、输入输出规格、平台差异910B / 310P以及从 Operation 到 Ops Runner 的完整调用链可用于快速定位源码、排查执行错误或作为理解其他 82 个 ATB 算子的范本。1. kv_cache 在 ATB 知识库中的定位在 ATB Agent 知识库中每个算子都有两份配套文档routing/{op_name}.md路由文件与ops/{category}/{op_name}/index.md知识条目使用方式为先查索引 → 再读路由 → 最后读知识条目见 .agent/knowledge/README.md。kv_cache算子的元信息如下来自 kv_cache 路由文件 与 主索引维度值分类infer推理复杂度 TierS简单算子4-8 文件单一路径文件数4Runner 类型OpsRunner, OperationACLNN 支持no纯原生 Ops 路径知识条目状态completeops/other/kv_cache/index.md主索引将kv_cache归入other功能分类纯 Ops 路径共 57 个算子之一同时标注了 Dtype 断点⚠️提示不同 dtype 下存在行为差异。2. 文件清单与角色划分路由文件给出了该算子的全部 4 个源码文件及其角色全部位于 src/ops/ops_infer/kv_cache/#文件角色1kv_cache_operation.cppOperation 定义2kv_cache_operation.hOperation 定义3kv_cache_ops_runner.cppOps Runner4kv_cache_ops_runner.hOps Runner这是 ATB 推理算子的标准目录模式src/ops/ops_infer/{op_name}/*_operation.*负责算子注册、参数校验与 InferShape*_ops_runner.*负责将算子映射为底层原生 Ops 图Mki 图并执行。因为该算子无 ACLNN 封装所以没有*_aclnn_runner.*文件文件数保持为 4。3. 推荐阅读顺序路由文件明确给出了由浅入深的阅读顺序这与 ATB 的运行时分层完全一致顺序文件重点关注1kv_cache_operation.h了解输入输出数量、InferShape 签名2kv_cache_operation.cppCreateRunner()决策逻辑3kv_cache_ops_runner.h原生 Ops 执行接口4kv_cache_ops_runner.cpp原生 Ops 调用链 平台适配下面按此顺序展开源码级解读。4. 第一站Operation 头文件——输入输出数量与接口签名kv_cache_operation.h 定义了KvCacheOperation : public OperationBase其公开接口揭示了算子的基本形状GetInputNum()返回5GetOutputNum()返回0InferShapeImpl()/InferShapeCheckImpl()/SetupCheckImpl()分别负责形状推断与运行时校验内部还声明了两套平台校验入口InferShapeDimCheck/SetupDimCheck910B与InferShapeDimCheck310P/SetupDimCheck310P310P。GetOutputNum() 0是该算子的关键特征KV Cache 写入是就地in-place更新输出直接写回输入past因此不需要额外分配输出张量。5. 第二站Operation 实现——注册、校验与 Runner 决策kv_cache_operation.cpp 是算子的核心决策层包含以下逻辑5.1 算子注册template Status CreateOperation(const infer::KvCacheParam opParam, Operation **operation)通过CreateOperation模板特化完成算子工厂注册OP_PARAM_RSV_CHECK(opParam)校验参数头文件的预留字段随后以std::nothrow方式构造KvCacheOperation实例。构造函数中通过GetSingletonAtbOperationIrCfg().GetOperationIr(KvCacheOperation)从算子 IR 配置中加载该算子的执行描述operationIr_。5.2 形状校验与平台分支InferShapeImpl仅记录日志后直接返回NO_ERROR——因为算子没有输出无需推导输出形状。真正的校验发生在InferShapeCheckImpl/SetupCheckImpl二者都通过GetSingletonConfig().Is910B()选择校验路径Status st GetSingletonConfig().Is910B() ? InferShapeDimCheck(inTensorDescs) : InferShapeDimCheck310P(inTensorDescs);910B 路径InferShapeDimCheck的约束包括输入 0newKv支持2 维或 4 维输入 1layerId、输入 3tokenOffset、输入 4seqLen均为 1 维输入 2past为 4 维past第 3 维为 hiddenSizefloat16 时必须是 16 的倍数int8 时必须是 32 的倍数内存对齐要求newKv为 2 维时其第 1 维、为 4 维时其后两维乘积均须等于 hiddenSizepast的 batch 维第 1 维必须与tokenOffset、seqLen的长度一致。310P 路径InferShapeDimCheck310P的约束不同newKv固定 4 维past变为5 维hiddenSize 取newKv第 1 维并与past第 2 维比对。5.3 Runner 决策std::shared_ptrRunner KvCacheOperation::CreateRunner(Context context) const { (void)context; return std::make_sharedKvCacheOpsRunner(param_); }CreateRunner()无条件返回KvCacheOpsRunner——由于该算子没有 ACLNN 路径唯一的执行通道就是原生 Ops Runner这正是路由文件标注 Runner 类型: OpsRunner, Operation 的源码依据。6. 第三站Ops Runner——Mki 图构建与原生 Ops 调用链kv_cache_ops_runner.h 中KvCacheOpsRunner : public OpsRunner的声明非常简洁仅持有infer::KvCacheParam param_成员。真正的逻辑集中在 kv_cache_ops_runner.cpp 的构造函数中6.1 输入张量映射const std::size_t intensorSize 5; kernelGraph_.inTensors.resize(intensorSize); kernelGraph_.outTensors.resize(0); Mki::Tensor newKvTensor kernelGraph_.inTensors.at(inTensorStart); Mki::Tensor layerIdTensor kernelGraph_.inTensors.at(inTensorStart); Mki::Tensor pastTensor kernelGraph_.inTensors.at(inTensorStart); Mki::Tensor tokenOffsetTensor kernelGraph_.inTensors.at(inTensorStart); Mki::Tensor seqLenTensor kernelGraph_.inTensors.at(inTensorStart);与 Operation 层一致Runner 构造 5 个输入张量、0 个输出张量。6.2 原生算子节点kernelGraph_.nodes.resize(1); auto kvCacheNode kernelGraph_.nodes.at(0); AtbOps::OpParam::KVCache kvCacheParam; kvCacheParam.type AtbOps::OpParam::KVCache::KVCACHE_ND; kvCacheParam.qSeqLen {}; kvCacheParam.kvSeqLen {}; kvCacheParam.batchRunStatus {}; kvCacheNode.opDesc {0, KVCacheOperation, kvCacheParam}; kvCacheNode.inTensors {newKvTensor, layerIdTensor, pastTensor, tokenOffsetTensor, seqLenTensor}; kvCacheNode.outTensors {pastTensor};该节点通过AtbOps::OpParam::KVCache参数描述原生算子行为type KVCACHE_ND表明使用 ND 格式qSeqLen、kvSeqLen、batchRunStatus等动态序列相关字段置空。值得注意的是outTensors {pastTensor}输出张量直接指向输入past从 Mki 图层面再次确认了就地更新语义。末尾的REG_RUNNER_TYPE(KvCacheOpsRunner)与REG_OP_PARAM(AtbOps::OpParam::KVCache)完成 Runner 类型与原生算子参数的注册。6.3 输入视图变换TensorViewFunckvCacheNode.inTensorViewFuncs[0] [](const Mki::SVectorint64_t oldDims, Mki::SVectorint64_t newDims) { if (oldDims.size() 2) { newDims oldDims; } else if (oldDims.size() 4) { newDims {oldDims.at(0) * oldDims.at(1), oldDims.at(2) * oldDims.at(3)}; } else { ATB_LOG(ERROR) kvcache operation intensor[0] only support dimNum 2 or 4.; } };第 0 个输入newKv在进入原生算子前会做视图重排4 维形状如[batch, kvHead, seq, headSize]被折叠为 2 维[batch*kvHead, seq*headSize]2 维形状则原样透传与 Operation 层支持 2 维或 4 维的校验完全呼应。7. 参数与输入输出规格7.1 参数头文件路由文件将参数头文件指向 include/atb/infer_op_params.h推理算子统一参数头文件。其中KvCacheParam当前仅包含预留字段struct KvCacheParam { uint8_t rsv[8] {0}; };即该算子当前版本不暴露业务侧可配置参数行为完全由输入张量与底层内核决定头文件中同区域的ReshapeAndCacheParam及其KvCacheCfg枚举K_CACHE_V_CACHE/K_CACHE_V_BYPASS/K_CACHE_V_CACHE_NZ属于相邻的 reshape_and_cache 系列算子可作为对比参考。7.2 输入规格来自 ops 配置ops_configs/atb_ops_info.ini 中[KvCacheOperation]段给出了官方的输入 dtype/format 规格输入名称dtype 支持format 支持input0newKvfloat16 / float16 / int8nd / fractal_nz / ndinput1layerIdint32ndinput2pastfloat16 / float16 / int8nd / fractal_nz / ndinput3tokenOffsetint32ndinput4seqLenint32nd三组 dtype/format 组合对应三种执行变体float16nd、float16fractal_nzNZ 分形格式与int8nd。这也解释了主索引中标注的 Dtype 断点int8 变体下 hiddenSize 对齐要求32 的倍数与 float1616 的倍数不同。8. 关联 Kernel 与周边资源路由文件标注的关联 Kernel 目录为 src/kernels/mixkernels/laser_attentionATB 知识库将 Laser Attention 相关的 KV Cache 融合内核集中在此。知识条目 .agent/knowledge/ops/other/kv_cache/index.md 还给出了两条补充元数据Pipeline 为单阶段Runner 标注为ops_runner / aclnn_runner路由文件与主索引标注为纯 OpsRunner实际以路由文件为准。如果需要观察该算子的调用方式与 golden 对比可参考仓库中的同族示例与测试KV Cache 相关操作在 example/op_demo/ 下有reshape_and_cache、paged_cache_load、paged_attention等可直接运行的 demo可帮助理解 KV Cache 写入/读取在真实推理图如 PagedAttention中的位置。9. 总结一条完整的算子调用链将以上源码证据串联KvCacheOperation的完整调用链为用户通过CreateOperation(KvCacheParam)创建KvCacheOperation工厂注册于 kv_cache_operation.cpp执行时先经InferShapeCheckImpl/SetupCheckImpl按Is910B()选择 910B 或 310P 的维度校验CreateRunner()创建KvCacheOpsRunnerkv_cache_operation.cppRunner 构造函数构建 5 输入 0 输出的 Mki 图节点 opDesc 为KVCacheOperation、参数类型KVCACHE_NDoutTensors直接复用past实现就地写入kv_cache_ops_runner.cpp底层原生内核按newKv、layerId、past、tokenOffset、seqLen五元组执行 KV Cache 写入。这套路由文件 → Operation → Ops Runner → 原生算子的阅读方法论同样适用于 ATB 中其余 81 个算子是理解 ascend-transformer-boost 推理算子体系的高效入口。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考