ascend-transformer-boost BlockCopyOperation 源码级解析KVCache Block 搬运算子的文件路由、实现与测试指南【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boostBlockCopyOperation 是 CANN ascend-transformer-boost 推理侧infer提供的高性能融合算子用于按 block 粒度在 KVCache 与 VCache 之间完成索引驱动的数据搬运是 PagedAttention、MLA 等长序列推理场景中 KV 重排与缓存整理的底层支撑。本文以仓库知识路由文档 .agent/knowledge/routing/block_copy.md 为骨架结合 src/ops/ops_infer/block_copy 与 src/kernels/mixkernels/blockcopy 的完整实现、ops_configs/atb_ops_info.ini 的配置声明以及 tests 下的测试用例讲解该算子的源码组织结构、推荐阅读路径、核心接口语义、数据流与执行原理读完即可掌握如何定位、阅读、调用与验证这一算子。1. 算子定位与整体概况路由文件在文件头给出了该算子的第一手分类信息分类: infer |复杂度: S |文件数: 4Runner 类型: OpsRunner,Operation |ACLNN: no预估阅读时间: 3-5 分钟结合源码可以确认这些元信息的含义分类 infer该算子位于推理侧算子目录 src/ops/ops_infer/block_copy属于src/ops/ops_infer的推理算子集合而非训练侧ops_train算子复杂度 Ssingle知识条目 .agent/knowledge/ops/other/block_copy/index.md 的 YAML 头将其标记为tier: S, type: single即单节点、单阶段的小型算子Runner 类型为 OpsRunner, Operation从源码看BlockCopyOperation 继承自OperationBase其CreateRunner()返回 BlockCopyOpsRunner而BlockCopyOpsRunner继承自OpsRunner见 src/atb/runner/ops_runner.hACLNN: no该算子不提供 aclnn 接口封装直接以 ATB Operation 形式使用。2. 文件清单与推荐阅读顺序路由文件列出了该算子的 4 个核心文件并给出了针对性的阅读顺序这是进入源码的最佳地图#文件角色1block_copy_operation.cppOperation 定义2block_copy_operation.hOperation 定义3block_copy_ops_runner.cppOps Runner4block_copy_ops_runner.hOps Runner顺序文件重点关注1block_copy_operation.h了解输入输出数量、InferShape 签名2block_copy_operation.cppCreateRunner() 决策逻辑3block_copy_ops_runner.h原生 Ops 执行接口4block_copy_ops_runner.cpp原生 Ops 调用链 平台适配阅读顺序背后的逻辑是自顶向下的分层理解先看 Operation 对外暴露的接口输入/输出数量、InferShape/Setup 校验签名再看 Operation 如何创建 RunnerCreateRunner 决策随后进入 Runner 层查看其如何承接 ATB 的 Tensor 数据最终落到原生 OpsMki 体系的 kernel 调用链与不同平台的适配逻辑。3. 源码路径速查路由文件给出了三处关键源码路径与仓库实际目录一一对应Op 目录: src/ops/ops_infer/block_copyKernel 目录: src/kernels/mixkernels/blockcopy参数头文件: include/atb/infer_op_params.h需要说明的是路由文件中的 Kernel 目录写作src/kernels/mixkernels/laser_attention但仓库中 block_copy 实际对应的 Kernel 实现位于 src/kernels/mixkernels/blockcopy该目录包含 operation、tiling、op_kernel 与 CMake 构建文件。阅读时以 src/kernels/mixkernels/blockcopy 为准它才是 BlockCopy 内核的完整实现所在。3.1 参数定义BlockCopyParaminclude/atb/infer_op_params.h 中定义了推理侧参数结构体//! //! \struct BlockCopyParam //! //! \brief 将KVCache里通过src indices指定的block数据copy到dst indices指定的block位置上。 //! struct BlockCopyParam { //! //! \brief 预留参数 //! uint8_t rsv[16] {0}; };该结构体注释直接点明了算子的语义将 KVCache 中由 src indices 指定的 block 数据复制到 dst indices 指定的 block 位置。BlockCopyParam目前仅含 16 字节预留字段rsv实际行为完全由 5 个输入 Tensor 的索引与形状描述驱动。对应的内核侧参数 src/kernels/include/atbops/params/blockcopy.h 定义了一个type枚举用于区分缓存格式struct BlockCopy { enum Type { BLOCK_COPY_CACHE_ND 0, BLOCK_COPY_CACHE_NZ 1 }; Type type BLOCK_COPY_CACHE_ND; };该枚举在 tiling 阶段被用于选择 ND910B 默认或 NZ310P fractal_nz格式的 tiling 分支详见下文第 6 节。4. Operation 层接口语义与参数校验4.1 输入输出数量block_copy_operation.cpp 定义了固定数量的输入输出uint32_t BlockCopyOperation::GetInputNum() const { const uint32_t inTensorNum 5; return inTensorNum; } uint32_t BlockCopyOperation::GetOutputNum() const { return 0; }输入为 5 个 Tensor输出为 0——这是因为该算子采用in-place 语义直接改写输入的 KCache/VCache不额外产出输出 Tensor。这一点在 Runner 的 kernel graph 构造中也有印证见第 5 节。5 个输入的角色在 block_copy_ops_runner.cpp 中命名清晰输入索引名称含义0kCacheK 缓存形状[block_count, block_size, num_heads, head_size]1vCacheV 缓存形状与 kCache 一致2srcBlockIndices源 block 索引一维 int32长度 源 block 个数3dstBlockIndices目标 block 索引一维 int32长度 目标 block 个数4cumSum源 block 到目标 block 的映射前缀和一维 int32与 srcBlockIndices 同长其中 cumSum 的语义是「多对一」映射的关键cumSum[i]表示第 i 个源 block 覆盖到第几个目标 block 位置前闭后开区间[cumSum[i-1], cumSum[i])内的 dst 索引都复制该源 block 的内容由此实现「一个源 block 复制到多个目标位置」的广播式复制。4.2 平台限制CreateOperation 是算子工厂入口其中对运行平台做了硬性限制if (!GetSingletonConfig().Is910B() !GetSingletonConfig().Is310P()) { ATB_LOG(ERROR) only support Atlas 800I A2 inference product and Atlas 300I Duo inference product; return ERROR_INVALID_PARAM; }即该算子仅支持 Atlas 800I A2Ascend910B与 Atlas 300I DuoAscend310P两类推理产品在其他平台上创建算子会直接返回ERROR_INVALID_PARAM。这一限制在测试用例中同样有对应验证见第 8 节block_copy_err_soc用例。4.3 InferShape 与 Setup 校验InferShapeCheckImplblock_copy_operation.cpp在形状推导阶段执行以下约束kCache 与 vCache 形状必须完全相等TensorShapeEqual且均为 4 维CACHE_DIM 4srcBlockIndices、dstBlockIndices 均为一维INDICES_DIM 1cumSum 形状必须与 srcBlockIndices 形状相等srcBlockIndices[0] 与 dstBlockIndices[0]元素个数均不能超过 blockCountkCache.shape.dims[0]。SetupCheckImplblock_copy_operation.cpp在 Setup 阶段做了更严格的运行时校验除重复上述形状约束外还增加了310P 平台kCache/vCache 的 dtype 必须是ACL_FLOAT16并通过SetupDimCheck310PL148-L180检查对齐约束——NZ 格式要求最后一维为 16NZBLOCKSIZE 16且 dims[2] 是 16 的倍数ND 格式要求前三维的乘积是 16 的倍数910B 平台kCache/vCache 不允许ACL_FORMAT_FRACTAL_NZ格式返回ERROR_INVALID_TENSOR_FORMAT。这些校验逻辑与 ops_configs 中的格式声明、以及测试用例中的错误用例如block_copy_310P_dim_err、block_copy_310P_dtype_err一一对应是理解「什么样的输入是合法的」的权威依据。4.4 CreateRunner 决策逻辑路由文件要求重点关注CreateRunner()的决策逻辑实现位于 block_copy_operation.cppstd::shared_ptrRunner BlockCopyOperation::CreateRunner(Context context) const { (void)context; return std::make_sharedBlockCopyOpsRunner(param_); }BlockCopy 的 Runner 决策非常简单直接无条件创建BlockCopyOpsRunner不存在多 Runner 分支选择。这也符合其「复杂度 S」的定位——单一 Runner 类型OpsRunner无 aclnn 路径。5. Runner 层KernelGraph 构建与原生 Ops 调用链5.1 OpsRunner 基类机制BlockCopyOpsRunner继承自OpsRunnersrc/atb/runner/ops_runner.h而OpsRunner是 ATB 中负责将高层 Tensor 数据编排为底层 Mki kernel 执行图KernelGraph的 Runner 基类。它承担了SetupImpl、ExecuteImpl、tiling buffer 管理、workspace 计算、kernel cache 等核心职责BlockCopyOpsRunner只需在构造函数中完成 KernelGraph 的初始化即可。5.2 KernelGraph 构造block_copy_ops_runner.cpp 的构造函数完成了关键的数据流绑定BlockCopyOpsRunner::BlockCopyOpsRunner(const infer::BlockCopyParam param) : OpsRunner(BlockCopyOpsRunner), param_(param) { ATB_LOG(INFO) BlockCopyOpsRunner::BlockCopyOpsRunner called; kernelGraph_.inTensors.resize(5); // dim:5 size_t inTensorId 0; Mki::Tensor kCache kernelGraph_.inTensors.at(inTensorId); Mki::Tensor vCache kernelGraph_.inTensors.at(inTensorId); Mki::Tensor srcBlockIndices kernelGraph_.inTensors.at(inTensorId); Mki::Tensor dstBlockIndices kernelGraph_.inTensors.at(inTensorId); Mki::Tensor cumSum kernelGraph_.inTensors.at(inTensorId); kernelGraph_.nodes.resize(1); auto blockCopyNode kernelGraph_.nodes.at(0); AtbOps::OpParam::BlockCopy blockCopyNodeParam {}; blockCopyNode.opDesc {0, BlockCopyOperation, blockCopyNodeParam}; blockCopyNode.inTensors {kCache, vCache, srcBlockIndices, dstBlockIndices, cumSum}; blockCopyNode.outTensors {kCache, vCache}; }这里可以清晰看到两层绑定关系ATB 侧 5 个输入 Tensor 与 Mki 侧 kernel 图输入的绑定kernelGraph_.inTensors依次挂接 kCache、vCache、srcBlockIndices、dstBlockIndices、cumSumkernel 节点的输入输出绑定节点的inTensors为 5 个输入outTensors直接复用kCache, vCache印证了第 4.1 节的结论——输出与输入共享同一块内存属于就地in-place更新。节点参数AtbOps::OpParam::BlockCopy blockCopyNodeParam {}使用默认构造type BLOCK_COPY_CACHE_ND随后由REG_RUNNER_TYPE(BlockCopyOpsRunner)与REG_OP_PARAM(AtbOps::OpParam::BlockCopy)两个宏完成 Runner 与参数类型的注册。5.3 原生 Ops 侧实现Mki 侧的 Operation 定义位于 src/kernels/mixkernels/blockcopy/blockcopy_operation.cpp它通过GetInputNum/GetOutputNum声明 5 入 2 出并在InferShapeImpl中直接将 outTensors 赋值为输入 K/V 缓存。其私有校验函数明确给出了各输入的数据类型与形状约束CheckKVCacheK/V 形状必须相同、均为 4 维dtype 支持float16、bf16、int8CheckSrcBlockListsrc 与 cumSum 均为一维int32且形状相同CheckDistBlockIndicesdst 为一维int32。GetBestKernel返回名为BlockCopyKernel的内核src/kernels/mixkernels/blockcopy/blockcopy_kernel.cpp由REG_OPERATION(BlockCopyOperation)注册。6. Kernel 层tiling 与多核并行搬运6.1 Tiling 数据与计算src/kernels/mixkernels/blockcopy/tiling/tiling_data.h 定义了下发到内核的 tiling 结构struct BlockCopyTilingData { uint32_t blockCount; // KVCache 中 block 总数 uint32_t blockSize; // 每个 block 的序列长度 uint32_t numHead; // 头数NZ 格式下固定为 1 uint32_t headSizeK; // K 的 head 维大小 uint32_t headSizeV; // V 的 head 维大小 uint32_t sourceCount; // 源 block 个数 uint32_t destinationCount;// 目标 block 个数 uint32_t typeByte; // 元素字节数 uint32_t blockDim; // 实际启用的核数 uint32_t perCoreCopyCount;// 每核平均搬运的 block 数 uint32_t tailCoreCopyCount; // 尾部核多搬运的 block 数 };blockcopy_tiling.cpp 的BlockCopyTiling是 tiling 主入口其核心逻辑包括按param.type选择 ND 或 NZ 分支BlockCopyTilingNd910BK/V 形状[block_count, block_size, num_heads, head_size]与BlockCopyTilingNz310PNZ 布局下无法直接取得头与头大小因此将numHead固定为 1headSize通过dims[1] * 16反推310P 平台额外执行 32 字节对齐检查BlockCopyTilingCheck310P多核并行任务切分actualCore min(destinationCount, vector 核数)perCoreCopyCount destinationCount / actualCoretailCoreCopyCount destinationCount % actualCore即把目标 block 的搬运任务按核均分多余的尾部任务由前若干核多承担一个blockDim与各核偏移随之确定以数据类型字节数生成 tilingKeyTILING_DTYPE_IDX * typeByte供内核选择对应 dtype 的模板实例。6.2 内核实现910B 与 310P构建文件 src/kernels/mixkernels/blockcopy/CMakeLists.txt 展示了双平台内核注册方式add_operation(BlockCopyOperation ${blockcopy_srcs}) add_kernel(blockcopy ascend910b vector op_kernel/blockcopy.cpp BlockCopyKernel) add_kernel(blockcopy ascend310p vector op_kernel/blockcopy_310p.cpp BlockCopyKernel)910B 与 310P 分别编译 op_kernel/blockcopy.cpp 与 op_kernel/blockcopy_310p.cpp两者都基于 AscendC 编写、注册为同一个BlockCopyKernel名。910B 版本内核入口extern C __global__ __aicore__ void blockcopy(GM_ADDR kCache, GM_ADDR vCache, GM_ADDR srcBlockIndices, GM_ADDR dstBlockIndices, GM_ADDR cumSum, GM_ADDR kCacheOut, GM_ADDR vCacheOut, GM_ADDR tiling)内核执行的核心流程Process()见 blockcopy.cpp分两个阶段Search 阶段定位源 block 偏移按 128 个元素为一组TILE_LENGTH 128分块加载 cumSum用CompareScalar Select ReduceSum向量指令二分式地逐 tile 比较找到当前核起始gmOffset对应的源 block 位置cumSumOffset_即通过前缀和反查「本核要搬运的第一个 block 是第几个源 block」Copy 阶段搬运数据将每个源 block 对应的目标 block 索引与 cumSum 对位加载在CopyOneSrc2MultiDst中按「一个源 block 复制到多个 dst 位置」的方式逐个执行DataCopyPadK/V 各一条队列src2dstQueueK/src2dstQueueV双缓冲BUFFER_NUM 2完成从kCacheGm[srcBlockIndex]到kCacheGm[dstBlockIndex]的搬运。由于 K/V 各自块大小blockSizeinElement_ blockSize * numHead * headSizeK可能超过单次 UB 容量OUT_UB_SIZE 45KB单 block 内部还按cacheCopyLoopCount_分片搬运确保任意规模的 block 都能正确复制。7. 配置声明atb_ops_info.iniops_configs/atb_ops_info.ini 中 [BlockCopyOperation] 段声明了算子输入输出的 dtype 与 format 组合该文件用于算子信息的统一登记与测试数据生成[BlockCopyOperation] input0.namekcache input0.dtypefloat16,bf16,int8,float16 input0.formatnd,nd,nd,fractal_nz input1.namevcache input1.dtypefloat16,bf16,int8,float16 input1.formatnd,nd,nd,fractal_nz input2.namesrcIndices input2.dtypeint32,int32,int32,int32 input2.formatnd,nd,nd,nd input3.namedstIndices input3.dtypeint32,int32,int32,int32 input3.formatnd,nd,nd,nd input4.namecumSum input4.dtypeint32,int32,int32,int32 input4.formatnd,nd,nd,nd output0.namekcacheOut output0.dtypefloat16,bf16,int8,float16 output0.formatnd,nd,nd,fractal_nz output1.namevcacheOut output1.dtypefloat16,bf16,int8,float16 output1.formatnd,nd,nd,fractal_nz该声明揭示的合法输入组合为 4 组nd / float16nd / bf16nd / int8fractal_nz / float16对应 310P 的 NZ 场景其中 srcIndices、dstIndices、cumSum 恒为nd / int32。这与 4.3 节的源码校验910B 不允许 NZ、310P 仅支持 fp16保持一致也解释了 ops 配置中第 4 组 fractal_nz 组合仅作用于 310P 平台。8. 测试验证从 CSV 用例到 Python 真值计算BlockCopy 在仓库中拥有完整的测试矩阵可作为理解算子行为与验证实现的样本。8.1 opstest CSV 用例tests/apitest/opstest/csv/block_copy.csv 覆盖了功能与异常两条线功能用例block_copy_fp16/bf16/int8910Bshape100,2,2,2src/dst/cumSum 均为 30 长度以及一组 310P 的block_copy_310P_fp16_nd/nz用例如64,30,16,16的 nd 与 fractal_nz 组合预期结果均为NO_ERROR异常用例block_copy_ini_errfloat dtype 非法预期ERROR_INVALID_TENSOR_INI_MATCH、block_copy_dim_diff_errK/V 形状不一致、block_copy_dim_errK/V 非 4 维、block_copy_dim_diff_err2cumSum 与 src 形状不一致、block_copy_index_dim_err索引长度超过 blockCount、block_copy_err_socAscend310B 平台预期ERROR_INVALID_PARAM、以及 310P 的 dtype/对齐错误用例与 4.3 节源码校验逐一对应。8.2 Python 真值实现tests/apitest/opstest/python/operations/block_copy/test_block_copy.py 给出了最直白的算子语义参考其 golden 计算逻辑def golden_calc(self, in_tensors): kCacheGolden self.kCache.copy() vCacheGolden self.vCache.copy() srcBlks in_tensors[2].cpu().numpy() dstBlks in_tensors[3].cpu().numpy() cumSum in_tensors[4].cpu().numpy() startIdx 0 endIdx 0 for i in range(len(srcBlks)): srcBlk srcBlks[i] endIdx cumSum[i] for j in range(startIdx, endIdx): dstBlk dstBlks[j] kCacheGolden[dstBlk] kCacheGolden[srcBlk].copy() vCacheGolden[dstBlk] vCacheGolden[srcBlk].copy() startIdx endIdx return [kCacheGolden, vCacheGolden]即对每个源 blocksrcBlks[i]将其复制到dstBlks[cumSum[i-1] : cumSum[i]]区间内的所有目标 block 位置测试通过np.array_equal逐元素比对算子输出与该真值。测试数据生成函数generate_dstBlks保证 dst 索引不落在 src 集合内避免源数据被覆盖导致的读取漂移generate_cumSum从[1, dstCount)中随机采样升序序列构造合法前缀和。内核级测试 tests/apitest/kernelstest/mix/test_blockcopy.py 与高层泛化测试 tests/high_level_test/BlockCopyOperation/Smoke/BlockCopyOperation_TestCase.csv200 组随机 shape 组合进一步扩大了 dtypefp16/bf16/int8、shapeblockCount 从 15 到 20000与平台910B/310P的覆盖范围其中blockCount20000, blockSize2, numHead2, headSize2这类大 block 数用例直接验证了第 6 节所述多核任务切分逻辑在大规模场景下的正确性。9. 快速导航与相关文档路由入口.agent/knowledge/routing/block_copy.md详细知识条目.agent/knowledge/ops/other/block_copy/index.md主索引infer 分类导航.agent/knowledge/README.md推理算子目录清单src/ops/ops_infer/AGENTS.md若希望从零开始跑通 BlockCopy 的功能验证可参考 tests/apitest/opstest/csv/block_copy.csv 中SocVersion为Ascend910B/Ascend310P的用例组织方式配合 opstest 框架在对应 SoC 环境执行算子本身不提供 aclnn 接口调用方式以 ATB Operation 为主。小结BlockCopyOperation 虽然只有 4 个源文件、被路由文件标记为「复杂度 S」但其背后是完整的 Operation → OpsRunner → KernelGraph → Mki Operation → Tiling → AscendC Kernel 六层链路ATB 侧负责形状/平台校验与就地绑定的 kernel graph 构建Mki 侧负责多对一索引映射的合法性检查tiling 负责跨核均分搬运任务内核则利用前缀和反查与双缓冲流水实现 K/V 双缓存的高效 block 级搬运。理解这一从路由文档到源码的完整脉络不仅能够快速上手 BlockCopy也为阅读 ascend-transformer-boost 中其他 infer 算子提供了可复用的方法。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考