尧图网站建设 
ascend-transformer-boost 中 PagedCacheLoad 算子导读路由文件、参数语义、形状校验与 Runner 执行链【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本篇基于仓库路由文件 paged_cache_load.md 展开完整还原 PagedCacheLoad 算子Paged KV Cache 加载即 ReshapeAndCache 的逆过程在 ascend-transformer-boost 中的 5 个核心文件清单、推荐阅读顺序、PagedCacheLoadParam参数语义、Operation 层的形状推导与校验规则以及 OpsRunner 构建 KernelGraph 的完整调用链帮助读者在不读全部源码的前提下精准定位并理解该算子的实现。1. 路由文件定位PagedCacheLoad 是什么.agent/knowledge/routing/paged_cache_load.md 是仓库知识库中为 PagedCacheLoad 算子编写的“路由文件”routing file它为开发者提供了一张精确的源码导航图。文件头部给出的元信息如下属性值含义分类infer属于推理inference类算子复杂度M中等复杂度文件数5算子由 5 个核心 C 文件构成Runner 类型OpsRunner, Operation采用OpsRunnerOperation双层架构ACLNNno未接入 ACLNN 接口使用 Operation/ACL C 接口从算子语义看PagedCacheLoad 服务于 Paged Attention 推理场景KV Cache 以“分页”block方式存储本算子根据blockTables页表和contextLens每个序列的实际长度把分页存储的 K/V 数据按序“抽取”成连续的 ND密集张量outKey/outValue供后续 Attention 计算使用。路由文件明确给出的三处源码位置是Op 目录src/ops/ops_infer/paged_cache_load/Kernel 目录src/kernels/mixkernels/paged_cache_load参数头文件include/atb/infer_op_params.h2. 文件清单路由文件第 1 节列出的 5 个文件及其角色如下路径均已转换为仓库根目录相对路径#文件角色1atb_acl_paged_cache_load.cppACL C 接口封装辅助文件2paged_cache_load_operation.cppOperation 定义实现3paged_cache_load_operation.hOperation 定义声明4paged_cache_load_ops_runner.cppOps Runner实现5paged_cache_load_ops_runner.hOps Runner声明对应的 Kernel 侧文件位于 src/kernels/mixkernels/paged_cache_load从目录结构看按 NZ/ND 两种 KV Cache 格式拆分了设备端实现op_kernel/paged_cache_load_nd.cpp、op_kernel/paged_cache_load_nz.cpp并配有独立的 tiling 代码tiling/paged_cache_load_tiling.cpp等Host 侧入口为paged_cache_load_kernel.cpp与paged_cache_load_operation.cpp。3. 推荐阅读顺序路由文件第 2 节给出的阅读顺序及其“重点关注”列如下本文第 47 节即按此顺序逐一展开顺序文件重点关注1paged_cache_load_operation.h了解输入输出数量、InferShape 签名2paged_cache_load_operation.cppCreateRunner() 决策逻辑3paged_cache_load_ops_runner.h原生 Ops 执行接口4paged_cache_load_ops_runner.cpp原生 Ops 调用链 平台适配5atb_acl_paged_cache_load.cpp辅助文件4. 参数定义infer::PagedCacheLoadParam算子的全部 Host 侧配置集中在 include/atb/infer_op_params.h 的PagedCacheLoadParam结构中共 3 个有效字段加 61 字节预留区struct PagedCacheLoadParam { enum KvCacheCfg : int8_t { K_CACHE_V_CACHE_NZ 0, //! 默认值key_cache/value_cache 为 NZ 格式 K_CACHE_V_CACHE_ND, //! key_cache/value_cache 为 ND 格式 }; KvCacheCfg kvCacheCfg K_CACHE_V_CACHE_NZ; //! 输入 SeqLens 为累加和模式第 n 个 batch 所需元素数为前 n 个 batch 的累加和 bool isSeqLensCumsumMode false; //! 提供 SeqStart为每个 batch 提供在 blockTable 中的初始位置类似 offset bool hasSeqStarts false; uint8_t rsv[61] {0}; };参数语义与源码行为一一对应kvCacheCfg决定 K/V Cache 张量的存储格式NZ 分片格式或 ND 密集格式进而决定输入张量个数与 IR 选择见第 5 节isSeqLensCumsumModecontextLens是否为“前缀和”形式。ND 模式下形状校验会按此把blockTables的 batch 维期望值改为contextLens[0] - 1见 paged_cache_load_operation.cpphasSeqStarts启用后第 7 个输入seqStarts生效指定每个 batch 在blockTables中的起始下标。重要限制在 CreateOperation 入口 中硬编码了约束——当kvCacheCfg K_CACHE_V_CACHE_NZ时hasSeqStarts与isSeqLensCumsumMode不允许同时为 true否则直接返回ERROR_INVALID_PARAM并打印 “foramt NZ donot support hasSeqStartstrue or isSeqLensCumsumModetrue!”。5. Operation 层输入输出数量与 InferShape 签名5.1 输入输出定义paged_cache_load_operation.h 声明了PagedCacheLoadOperation : public OperationBase对外暴露的接口为GetInputNum()/GetOutputNum()输入输出张量数量InferShapeImpl()/InferShapeCheckImpl()形状推导与推导期校验CreateRunner()构建执行器SetupCheckImpl()Setup 期张量校验三个私有校验函数KVCacheDimCheck910BND、KVCacheDimCheck910BNZ、DimCheck。输入张量的索引常量定义在 paged_cache_load_operation.cpp索引名称语义0keyCache分页 K 缓存维度 41valueCache分页 V 缓存维度 42blockTables页表维度 2[batch, maxBlocksPerBatch]3contextLens每个序列实际长度维度 14key待写入输出位置的 K兼作输出模板NZ 下维度 2、ND 下维度 35value待写入输出位置的 V兼作输出模板6seqStarts每个 batch 在 blockTable 中的起始位置仅 ND 且hasSeqStarts时使用维度 1输入数量由kvCacheCfg动态决定GetInputNum()NZ 模式 6 个ND 模式 7 个输出恒为 2 个outKey、outValue。5.2 InferShape 签名InferShapeImpl的实现非常直接输出形状直接复制输入 4/5key/value的 TensorDesc。这意味着key/value输入张量同时充当输出张量的形状模板——加载结果被写回与key/value同形状的 ND 张量体现了“加载即覆盖”的原地语义。5.3 算子创建时的平台与配置校验路由文件提示第 2 个文件关注 “CreateRunner() 决策逻辑”实际决策链条从CreateOperation自由函数开始paged_cache_load_operation.cpp平台校验GetSingletonConfig().Is910B()为 false 时直接报错 “paged cache load only support Atlas 800I A2 inference product!”——该算子仅在 910BAtlas 800I A2 推理产品上可用配置校验kvCacheCfg必须落在K_CACHE_V_CACHE_NZ/K_CACHE_V_CACHE_ND区间内组合校验NZ (hasSeqStarts或isSeqLensCumsumMode) 组合被禁止。构造函数中还存在一处按配置分流L82-L88NZ 模式绑定 IR 名PagedCacheLoadOperationND 模式绑定PagedCacheLoadOperationND即同一算子在 Kernel IR 配置层面按格式注册了两套描述。而CreateRunner()本身没有分支——恒定返回PagedCacheLoadOpsRunner真正的 NZ/ND 决策被推迟到了 Runner 构造函数内部见第 7 节。6. Setup 期校验规则详解DimCheckL149-L209是所有张量校验的公共入口随后按格式分派到KVCacheDimCheck910BNZ或KVCacheDimCheck910BND。核心规则汇总如下公共规则规则说明numBlocks一致keyCache 与 valueCache 的dims[0]block 总数必须相等sumcontext一致key 与 value 的dims[0]总 token 数必须相等blockSize ≠ 0任一格式下 block 大小不可为 0batch 维匹配blockTables.dims[0]必须等于contextLens.dims[0]ND 模式且isSeqLensCumsumModetrue时等于contextLens.dims[0] - 1NZ 格式专属规则KVCacheDimCheck910BNZkeyCache/valueCache 为 4 维blockTables 2 维contextLens 1 维key/value 2 维blockSizekeyCache.dims[2]必须 16 字节对齐末维约束int8 时 keyCache 末维必须为 32valueCache 为 32同为 int8或 16非 int8非 int8 时两者末维均必须为 16容量上限numHeads × headSize按 32 字节/元素计不得超过 147456 字节即常量MAX_k/MAX_v 147456。ND 格式专属规则KVCacheDimCheck910BNDkeyCache/valueCache 为 4 维key/value 为 3 维seqStarts 1 维对齐约束int8 时num_heads × head_size必须 32 对齐fp16/bf16 时必须 16 对齐hasSeqStartstrue时seqStarts.dims[0]必须等于blockTables.dims[0]。此外SetupCheckImplL120-L147在真实张量就位后再次核对“输出形状与输入缓存形状对齐”NZ 下out.dims[1] cache.dims[1] × cache.dims[3]即 head 块数 × 32 对齐元素数 numHeads × headSizeND 下out.dims[1..2] cache.dims[2..3]。7. Runner 层KernelGraph 构建与平台适配7.1 接口声明paged_cache_load_ops_runner.h 中PagedCacheLoadOpsRunner继承自OpsRunner额外提供SetPaParam(AtbOps::OpParam::PagedCacheLoad)用于更新 kernel 参数并持有Mki::Tensor nullTensor_空 tensor 占位。7.2 调用链参数 → KernelGraph 节点paged_cache_load_ops_runner.cpp 的构造函数完成全部平台适配流程为预分配kernelGraph_.inTensors7 个与outTensors2 个按序绑定 7 个输入 tensorkeyCache、valueCache、blockTables、contextLens、key、value、seq_starts依据kvCacheCfg分两支构建唯一的 kernel 节点PagedCacheLoadOperation配置项NZ 分支ND 分支节点参数typePAGED_CACHE_LOAD_NZPAGED_CACHE_LOAD_NDcuSeqLens透传param_.isSeqLensCumsumMode同左hasSeqStarts透传param_.hasSeqStarts同左输入格式inferShapePreFunc输入 0/1cache设TENSOR_FORMAT_FRACTAL_NZ其余ND全部TENSOR_FORMAT_ND一个值得注意的细节NZ 分支复用了inTensors.at(3)contextLens 槽位作为seq_starts的引用——因为 NZ 模式本身不支持hasSeqStarts该槽位以占位方式挂入节点输入保证节点输入向量统一为 7 项L33-L59。文件末尾通过REG_RUNNER_TYPE(PagedCacheLoadOpsRunner)与REG_OP_PARAM(AtbOps::OpParam::PagedCacheLoad)完成运行器类型与算子参数的注册L87-L88使内核图调度框架可按名字找回执行器。8. ACL C 接口封装按路由文件建议的第 5 步atb_acl_paged_cache_load.cpp 提供 C 语言入口面向不便直接操作atb::Operation的调用方AtbPagedCacheLoadGetWorkspaceSizeL21-L78接收 7 个aclTensor*输入指针与kvCacheCfg、isSeqLensCumsumType、hasSeqStarts三个标量参数内部组装PagedCacheLoadParam并调用CreateOperation随后把 67 个输入、2 个输出直接复用 key/value 指针体现原地语义装入VariantPack并调用Setup求出 workspace 大小AtbPagedCacheLoadL81-L89持 workspace 调用op-Execute完成执行。注意该文件常量g_PAGED_CACHE_LOAD_INTENSOR_NUM 6只有hasSeqStartstrueND 模式时inTensors才扩容到 7与第 5 节GetInputNum()的规则一致。9. 结合示例 Demo 复现调用流程仓库提供了可直接编译运行的示例 example/op_demo/paged_cache_load/paged_cache_load_demo.cpp配套说明见同目录 README.md示例采用NZ int8配置关键常量常量值含义NUM_BLOCKS4分页 block 总数BLOCK_SIZE128每 block 的 token 数NUM_HEADS4注意力头数HEAD_SIZE_K/HEAD_SIZE_V32 / 32K/V 每头维度ELENUM_ALIGNED32NZ 末维对齐宽度int8 下 32LEN_CONTEXT_LENS3batch 大小NUM_TOKENS384输出总 token 数3 × 128示例构造流程对应PrepareOperationPrepareInTensoratb::infer::PagedCacheLoadParam opParam; opParam.kvCacheCfg atb::infer::PagedCacheLoadParam::K_CACHE_V_CACHE_NZ; opParam.isSeqLensCumsumMode false; opParam.hasSeqStarts false; atb::CreateOperation(opParam, op); // 910B 平台校验通过后创建 // K/V cacheNZ 4D{4 blocks, 4 head块, 128 blockSize, 32 对齐宽} // 形状 {NUM_BLOCKS, NUM_HEADS*HEAD_SIZE/ELENUM_ALIGNED, BLOCK_SIZE, ELENUM_ALIGNED} // blockTable {3,1} {2,0,0}contextLens {3} {128,128,128} // 输出 k/vND 2D {NUM_TOKENS, NUM_HEADS*HEAD_SIZE}主流程为标准 ATB 生命周期aclInit→atb::CreateContext→CreateOperation→ 填充VariantPackin 6 张量 / out 2 张量→Setup求 workspace→Execute→aclrtSynchronizeStream→ 释放设备内存并DestroyOperation。一个小细节示例中算子指针变量名为ropeOp系从 RoPE 示例模板沿用而来阅读源码时不要误以为本 demo 与 RoPE 相关。10. 约束与适用前提小结约束项取值/限制依据平台仅 910BAtlas 800I A2 推理产品Is910B()检查operation.cpp L57-L60KV 格式NZ 或 ND 二选一KvCacheCfg枚举NZ 模式不支持hasSeqStarts/isSeqLensCumsumModeCreateOperation校验NZ blockSize16 字节对齐DimCheckNZ 末维int8K32V32同为 int8或 16K int8 / V 非 int8非 int8均为 16KVCacheDimCheck910BNZND 对齐int8numHeads×headSize32 对齐fp16/bf1616 对齐KVCacheDimCheck910BNDNZ 容量上限numHeads × headSize 147456 字节常量MAX_k/MAX_v输入数NZ6ND7含 seqStartsGetInputNum()输出语义outKey/outValue 与输入 key/value 同形状原地加载InferShapeImpl ACL 接口复用指针如需进一步深入可沿路由文件给出的三条线索继续Op 层逻辑看 src/ops/ops_infer/paged_cache_load/ 下 5 个文件设备端实现与 tiling 看 src/kernels/mixkernels/paged_cache_load其中op_kernel/下按paged_cache_load_nd.cpp/paged_cache_load_nz.cpp区分两种格式接口契约以 include/atb/infer_op_params.h 中的PagedCacheLoadParam为准。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ascend-transformer-boost 中 PagedCacheLoad 算子导读：路由文件、参数语义、形状校验与 Runner 执行链
ascend-transformer-boost 中 PagedCacheLoad 算子导读路由文件、参数语义、形状校验与 Runner 执行链【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本篇基于仓库路由文件 paged_cache_load.md 展开完整还原 PagedCacheLoad 算子Paged KV Cache 加载即 ReshapeAndCache 的逆过程在 ascend-transformer-boost 中的 5 个核心文件清单、推荐阅读顺序、PagedCacheLoadParam参数语义、Operation 层的形状推导与校验规则以及 OpsRunner 构建 KernelGraph 的完整调用链帮助读者在不读全部源码的前提下精准定位并理解该算子的实现。1. 路由文件定位PagedCacheLoad 是什么.agent/knowledge/routing/paged_cache_load.md 是仓库知识库中为 PagedCacheLoad 算子编写的“路由文件”routing file它为开发者提供了一张精确的源码导航图。文件头部给出的元信息如下属性值含义分类infer属于推理inference类算子复杂度M中等复杂度文件数5算子由 5 个核心 C 文件构成Runner 类型OpsRunner, Operation采用OpsRunnerOperation双层架构ACLNNno未接入 ACLNN 接口使用 Operation/ACL C 接口从算子语义看PagedCacheLoad 服务于 Paged Attention 推理场景KV Cache 以“分页”block方式存储本算子根据blockTables页表和contextLens每个序列的实际长度把分页存储的 K/V 数据按序“抽取”成连续的 ND密集张量outKey/outValue供后续 Attention 计算使用。路由文件明确给出的三处源码位置是Op 目录src/ops/ops_infer/paged_cache_load/Kernel 目录src/kernels/mixkernels/paged_cache_load参数头文件include/atb/infer_op_params.h2. 文件清单路由文件第 1 节列出的 5 个文件及其角色如下路径均已转换为仓库根目录相对路径#文件角色1atb_acl_paged_cache_load.cppACL C 接口封装辅助文件2paged_cache_load_operation.cppOperation 定义实现3paged_cache_load_operation.hOperation 定义声明4paged_cache_load_ops_runner.cppOps Runner实现5paged_cache_load_ops_runner.hOps Runner声明对应的 Kernel 侧文件位于 src/kernels/mixkernels/paged_cache_load从目录结构看按 NZ/ND 两种 KV Cache 格式拆分了设备端实现op_kernel/paged_cache_load_nd.cpp、op_kernel/paged_cache_load_nz.cpp并配有独立的 tiling 代码tiling/paged_cache_load_tiling.cpp等Host 侧入口为paged_cache_load_kernel.cpp与paged_cache_load_operation.cpp。3. 推荐阅读顺序路由文件第 2 节给出的阅读顺序及其“重点关注”列如下本文第 47 节即按此顺序逐一展开顺序文件重点关注1paged_cache_load_operation.h了解输入输出数量、InferShape 签名2paged_cache_load_operation.cppCreateRunner() 决策逻辑3paged_cache_load_ops_runner.h原生 Ops 执行接口4paged_cache_load_ops_runner.cpp原生 Ops 调用链 平台适配5atb_acl_paged_cache_load.cpp辅助文件4. 参数定义infer::PagedCacheLoadParam算子的全部 Host 侧配置集中在 include/atb/infer_op_params.h 的PagedCacheLoadParam结构中共 3 个有效字段加 61 字节预留区struct PagedCacheLoadParam { enum KvCacheCfg : int8_t { K_CACHE_V_CACHE_NZ 0, //! 默认值key_cache/value_cache 为 NZ 格式 K_CACHE_V_CACHE_ND, //! key_cache/value_cache 为 ND 格式 }; KvCacheCfg kvCacheCfg K_CACHE_V_CACHE_NZ; //! 输入 SeqLens 为累加和模式第 n 个 batch 所需元素数为前 n 个 batch 的累加和 bool isSeqLensCumsumMode false; //! 提供 SeqStart为每个 batch 提供在 blockTable 中的初始位置类似 offset bool hasSeqStarts false; uint8_t rsv[61] {0}; };参数语义与源码行为一一对应kvCacheCfg决定 K/V Cache 张量的存储格式NZ 分片格式或 ND 密集格式进而决定输入张量个数与 IR 选择见第 5 节isSeqLensCumsumModecontextLens是否为“前缀和”形式。ND 模式下形状校验会按此把blockTables的 batch 维期望值改为contextLens[0] - 1见 paged_cache_load_operation.cpphasSeqStarts启用后第 7 个输入seqStarts生效指定每个 batch 在blockTables中的起始下标。重要限制在 CreateOperation 入口 中硬编码了约束——当kvCacheCfg K_CACHE_V_CACHE_NZ时hasSeqStarts与isSeqLensCumsumMode不允许同时为 true否则直接返回ERROR_INVALID_PARAM并打印 “foramt NZ donot support hasSeqStartstrue or isSeqLensCumsumModetrue!”。5. Operation 层输入输出数量与 InferShape 签名5.1 输入输出定义paged_cache_load_operation.h 声明了PagedCacheLoadOperation : public OperationBase对外暴露的接口为GetInputNum()/GetOutputNum()输入输出张量数量InferShapeImpl()/InferShapeCheckImpl()形状推导与推导期校验CreateRunner()构建执行器SetupCheckImpl()Setup 期张量校验三个私有校验函数KVCacheDimCheck910BND、KVCacheDimCheck910BNZ、DimCheck。输入张量的索引常量定义在 paged_cache_load_operation.cpp索引名称语义0keyCache分页 K 缓存维度 41valueCache分页 V 缓存维度 42blockTables页表维度 2[batch, maxBlocksPerBatch]3contextLens每个序列实际长度维度 14key待写入输出位置的 K兼作输出模板NZ 下维度 2、ND 下维度 35value待写入输出位置的 V兼作输出模板6seqStarts每个 batch 在 blockTable 中的起始位置仅 ND 且hasSeqStarts时使用维度 1输入数量由kvCacheCfg动态决定GetInputNum()NZ 模式 6 个ND 模式 7 个输出恒为 2 个outKey、outValue。5.2 InferShape 签名InferShapeImpl的实现非常直接输出形状直接复制输入 4/5key/value的 TensorDesc。这意味着key/value输入张量同时充当输出张量的形状模板——加载结果被写回与key/value同形状的 ND 张量体现了“加载即覆盖”的原地语义。5.3 算子创建时的平台与配置校验路由文件提示第 2 个文件关注 “CreateRunner() 决策逻辑”实际决策链条从CreateOperation自由函数开始paged_cache_load_operation.cpp平台校验GetSingletonConfig().Is910B()为 false 时直接报错 “paged cache load only support Atlas 800I A2 inference product!”——该算子仅在 910BAtlas 800I A2 推理产品上可用配置校验kvCacheCfg必须落在K_CACHE_V_CACHE_NZ/K_CACHE_V_CACHE_ND区间内组合校验NZ (hasSeqStarts或isSeqLensCumsumMode) 组合被禁止。构造函数中还存在一处按配置分流L82-L88NZ 模式绑定 IR 名PagedCacheLoadOperationND 模式绑定PagedCacheLoadOperationND即同一算子在 Kernel IR 配置层面按格式注册了两套描述。而CreateRunner()本身没有分支——恒定返回PagedCacheLoadOpsRunner真正的 NZ/ND 决策被推迟到了 Runner 构造函数内部见第 7 节。6. Setup 期校验规则详解DimCheckL149-L209是所有张量校验的公共入口随后按格式分派到KVCacheDimCheck910BNZ或KVCacheDimCheck910BND。核心规则汇总如下公共规则规则说明numBlocks一致keyCache 与 valueCache 的dims[0]block 总数必须相等sumcontext一致key 与 value 的dims[0]总 token 数必须相等blockSize ≠ 0任一格式下 block 大小不可为 0batch 维匹配blockTables.dims[0]必须等于contextLens.dims[0]ND 模式且isSeqLensCumsumModetrue时等于contextLens.dims[0] - 1NZ 格式专属规则KVCacheDimCheck910BNZkeyCache/valueCache 为 4 维blockTables 2 维contextLens 1 维key/value 2 维blockSizekeyCache.dims[2]必须 16 字节对齐末维约束int8 时 keyCache 末维必须为 32valueCache 为 32同为 int8或 16非 int8非 int8 时两者末维均必须为 16容量上限numHeads × headSize按 32 字节/元素计不得超过 147456 字节即常量MAX_k/MAX_v 147456。ND 格式专属规则KVCacheDimCheck910BNDkeyCache/valueCache 为 4 维key/value 为 3 维seqStarts 1 维对齐约束int8 时num_heads × head_size必须 32 对齐fp16/bf16 时必须 16 对齐hasSeqStartstrue时seqStarts.dims[0]必须等于blockTables.dims[0]。此外SetupCheckImplL120-L147在真实张量就位后再次核对“输出形状与输入缓存形状对齐”NZ 下out.dims[1] cache.dims[1] × cache.dims[3]即 head 块数 × 32 对齐元素数 numHeads × headSizeND 下out.dims[1..2] cache.dims[2..3]。7. Runner 层KernelGraph 构建与平台适配7.1 接口声明paged_cache_load_ops_runner.h 中PagedCacheLoadOpsRunner继承自OpsRunner额外提供SetPaParam(AtbOps::OpParam::PagedCacheLoad)用于更新 kernel 参数并持有Mki::Tensor nullTensor_空 tensor 占位。7.2 调用链参数 → KernelGraph 节点paged_cache_load_ops_runner.cpp 的构造函数完成全部平台适配流程为预分配kernelGraph_.inTensors7 个与outTensors2 个按序绑定 7 个输入 tensorkeyCache、valueCache、blockTables、contextLens、key、value、seq_starts依据kvCacheCfg分两支构建唯一的 kernel 节点PagedCacheLoadOperation配置项NZ 分支ND 分支节点参数typePAGED_CACHE_LOAD_NZPAGED_CACHE_LOAD_NDcuSeqLens透传param_.isSeqLensCumsumMode同左hasSeqStarts透传param_.hasSeqStarts同左输入格式inferShapePreFunc输入 0/1cache设TENSOR_FORMAT_FRACTAL_NZ其余ND全部TENSOR_FORMAT_ND一个值得注意的细节NZ 分支复用了inTensors.at(3)contextLens 槽位作为seq_starts的引用——因为 NZ 模式本身不支持hasSeqStarts该槽位以占位方式挂入节点输入保证节点输入向量统一为 7 项L33-L59。文件末尾通过REG_RUNNER_TYPE(PagedCacheLoadOpsRunner)与REG_OP_PARAM(AtbOps::OpParam::PagedCacheLoad)完成运行器类型与算子参数的注册L87-L88使内核图调度框架可按名字找回执行器。8. ACL C 接口封装按路由文件建议的第 5 步atb_acl_paged_cache_load.cpp 提供 C 语言入口面向不便直接操作atb::Operation的调用方AtbPagedCacheLoadGetWorkspaceSizeL21-L78接收 7 个aclTensor*输入指针与kvCacheCfg、isSeqLensCumsumType、hasSeqStarts三个标量参数内部组装PagedCacheLoadParam并调用CreateOperation随后把 67 个输入、2 个输出直接复用 key/value 指针体现原地语义装入VariantPack并调用Setup求出 workspace 大小AtbPagedCacheLoadL81-L89持 workspace 调用op-Execute完成执行。注意该文件常量g_PAGED_CACHE_LOAD_INTENSOR_NUM 6只有hasSeqStartstrueND 模式时inTensors才扩容到 7与第 5 节GetInputNum()的规则一致。9. 结合示例 Demo 复现调用流程仓库提供了可直接编译运行的示例 example/op_demo/paged_cache_load/paged_cache_load_demo.cpp配套说明见同目录 README.md示例采用NZ int8配置关键常量常量值含义NUM_BLOCKS4分页 block 总数BLOCK_SIZE128每 block 的 token 数NUM_HEADS4注意力头数HEAD_SIZE_K/HEAD_SIZE_V32 / 32K/V 每头维度ELENUM_ALIGNED32NZ 末维对齐宽度int8 下 32LEN_CONTEXT_LENS3batch 大小NUM_TOKENS384输出总 token 数3 × 128示例构造流程对应PrepareOperationPrepareInTensoratb::infer::PagedCacheLoadParam opParam; opParam.kvCacheCfg atb::infer::PagedCacheLoadParam::K_CACHE_V_CACHE_NZ; opParam.isSeqLensCumsumMode false; opParam.hasSeqStarts false; atb::CreateOperation(opParam, op); // 910B 平台校验通过后创建 // K/V cacheNZ 4D{4 blocks, 4 head块, 128 blockSize, 32 对齐宽} // 形状 {NUM_BLOCKS, NUM_HEADS*HEAD_SIZE/ELENUM_ALIGNED, BLOCK_SIZE, ELENUM_ALIGNED} // blockTable {3,1} {2,0,0}contextLens {3} {128,128,128} // 输出 k/vND 2D {NUM_TOKENS, NUM_HEADS*HEAD_SIZE}主流程为标准 ATB 生命周期aclInit→atb::CreateContext→CreateOperation→ 填充VariantPackin 6 张量 / out 2 张量→Setup求 workspace→Execute→aclrtSynchronizeStream→ 释放设备内存并DestroyOperation。一个小细节示例中算子指针变量名为ropeOp系从 RoPE 示例模板沿用而来阅读源码时不要误以为本 demo 与 RoPE 相关。10. 约束与适用前提小结约束项取值/限制依据平台仅 910BAtlas 800I A2 推理产品Is910B()检查operation.cpp L57-L60KV 格式NZ 或 ND 二选一KvCacheCfg枚举NZ 模式不支持hasSeqStarts/isSeqLensCumsumModeCreateOperation校验NZ blockSize16 字节对齐DimCheckNZ 末维int8K32V32同为 int8或 16K int8 / V 非 int8非 int8均为 16KVCacheDimCheck910BNZND 对齐int8numHeads×headSize32 对齐fp16/bf1616 对齐KVCacheDimCheck910BNDNZ 容量上限numHeads × headSize 147456 字节常量MAX_k/MAX_v输入数NZ6ND7含 seqStartsGetInputNum()输出语义outKey/outValue 与输入 key/value 同形状原地加载InferShapeImpl ACL 接口复用指针如需进一步深入可沿路由文件给出的三条线索继续Op 层逻辑看 src/ops/ops_infer/paged_cache_load/ 下 5 个文件设备端实现与 tiling 看 src/kernels/mixkernels/paged_cache_load其中op_kernel/下按paged_cache_load_nd.cpp/paged_cache_load_nz.cpp区分两种格式接口契约以 include/atb/infer_op_params.h 中的PagedCacheLoadParam为准。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考