企业官网建设流程全解析

ascend-transformer-boost BlockCopyOperation 源码级解析：KVCache Block 搬运算子的文件路由、实现与测试指南

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

BlockCopyOperation 是 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）算子；
• 复杂度 S（single）：知识条目 .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.h）；
• ACLNN: no：该算子不提供 aclnn 接口封装，直接以 ATB Operation 形式使用。

2. 文件清单与推荐阅读顺序

路由文件列出了该算子的 4 个核心文件，并给出了针对性的阅读顺序，这是进入源码的最佳地图：

阅读顺序背后的逻辑是自顶向下的分层理解：先看 Operation 对外暴露的接口（输入/输出数量、InferShape/Setup 校验签名），再看 Operation 如何创建 Runner（CreateRunner 决策），随后进入 Runner 层查看其如何承接 ATB 的 Tensor 数据，最终落到原生 Ops（Mki 体系）的 kernel 调用链与不同平台的适配逻辑。

3. 源码路径速查

路由文件给出了三处关键源码路径，与仓库实际目录一一对应：

• Op 目录: src/ops/ops_infer/block_copy
• Kernel 目录: 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 参数定义（BlockCopyParam）

include/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 阶段被用于选择 ND（910B 默认）或 NZ（310P 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 中命名清晰：

其中 cumSum 的语义是「多对一」映射的关键：cumSum[i]表示第 i 个源 block 覆盖到第几个目标 block 位置（前闭后开区间[cumSum[i-1], cumSum[i])内的 dst 索引都复制该源 block 的内容），由此实现「一个源 block 复制到多个目标位置」的广播式复制。

4.2 平台限制

CreateOperation 是算子工厂入口，其中对运行平台做了硬性限制：

if (!GetSingleton<Config>().Is910B() && !GetSingleton<Config>().Is310P()) { ATB_LOG(ERROR) << "only support Atlas 800I A2 inference product and Atlas 300I Duo inference product"; return ERROR_INVALID_PARAM; }

即该算子仅支持 Atlas 800I A2（Ascend910B）与 Atlas 300I Duo（Ascend310P）两类推理产品，在其他平台上创建算子会直接返回ERROR_INVALID_PARAM。这一限制在测试用例中同样有对应验证（见第 8 节block_copy_err_soc用例）。

4.3 InferShape 与 Setup 校验

InferShapeCheckImpl（block_copy_operation.cpp）在形状推导阶段执行以下约束：

• kCache 与 vCache 形状必须完全相等（TensorShapeEqual），且均为 4 维（CACHE_DIM = 4）；
• srcBlockIndices、dstBlockIndices 均为一维（INDICES_DIM = 1）；
• cumSum 形状必须与 srcBlockIndices 形状相等；
• srcBlockIndices[0] 与 dstBlockIndices[0]（元素个数）均不能超过 blockCount（kCache.shape.dims[0]）。

SetupCheckImpl（block_copy_operation.cpp）在 Setup 阶段做了更严格的运行时校验，除重复上述形状约束外，还增加了：

• 310P 平台：kCache/vCache 的 dtype 必须是ACL_FLOAT16，并通过SetupDimCheck310P（L148-L180）检查对齐约束——NZ 格式要求最后一维为 16（NZBLOCKSIZE = 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.cpp：

std::shared_ptr<Runner> BlockCopyOperation::CreateRunner(Context &context) const { (void)context; return std::make_shared<BlockCopyOpsRunner>(param_); }

BlockCopy 的 Runner 决策非常简单直接：无条件创建BlockCopyOpsRunner，不存在多 Runner 分支选择。这也符合其「复杂度 S」的定位——单一 Runner 类型（OpsRunner），无 aclnn 路径。

5. Runner 层：KernelGraph 构建与原生 Ops 调用链

5.1 OpsRunner 基类机制

BlockCopyOpsRunner继承自OpsRunner（src/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}; }

这里可以清晰看到两层绑定关系：

1. ATB 侧 5 个输入 Tensor 与 Mki 侧 kernel 图输入的绑定：kernelGraph_.inTensors依次挂接 kCache、vCache、srcBlockIndices、dstBlockIndices、cumSum；
2. kernel 节点的输入输出绑定：节点的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 缓存。其私有校验函数明确给出了各输入的数据类型与形状约束：

• CheckKVCache：K/V 形状必须相同、均为 4 维，dtype 支持float16、bf16、int8；
• CheckSrcBlockList：src 与 cumSum 均为一维int32且形状相同；
• CheckDistBlockIndices：dst 为一维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 分支：BlockCopyTilingNd（910B，K/V 形状[block_count, block_size, num_heads, head_size]）与BlockCopyTilingNz（310P，NZ 布局下无法直接取得头与头大小，因此将numHead固定为 1，headSize通过dims[1] * 16反推）；
• 310P 平台额外执行 32 字节对齐检查（BlockCopyTilingCheck310P）；
• 多核并行任务切分：actualCore = min(destinationCount, vector 核数)，perCoreCopyCount = destinationCount / actualCore，tailCoreCopyCount = destinationCount % actualCore，即把目标 block 的搬运任务按核均分，多余的尾部任务由前若干核多承担一个，blockDim与各核偏移随之确定；
• 以数据类型字节数生成 tilingKey（TILING_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）分两个阶段：

1. Search 阶段（定位源 block 偏移）：按 128 个元素为一组（TILE_LENGTH = 128）分块加载 cumSum，用CompareScalar + Select + ReduceSum向量指令二分式地（逐 tile 比较）找到当前核起始gmOffset对应的源 block 位置cumSumOffset_，即通过前缀和反查「本核要搬运的第一个 block 是第几个源 block」；
2. Copy 阶段（搬运数据）：将每个源 block 对应的目标 block 索引与 cumSum 对位加载，在CopyOneSrc2MultiDst中按「一个源 block 复制到多个 dst 位置」的方式逐个执行DataCopyPad（K/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.ini

ops_configs/atb_ops_info.ini 中 [BlockCopyOperation] 段声明了算子输入输出的 dtype 与 format 组合（该文件用于算子信息的统一登记与测试数据生成）：

[BlockCopyOperation] input0.name=kcache input0.dtype=float16,bf16,int8,float16 input0.format=nd,nd,nd,fractal_nz input1.name=vcache input1.dtype=float16,bf16,int8,float16 input1.format=nd,nd,nd,fractal_nz input2.name=srcIndices input2.dtype=int32,int32,int32,int32 input2.format=nd,nd,nd,nd input3.name=dstIndices input3.dtype=int32,int32,int32,int32 input3.format=nd,nd,nd,nd input4.name=cumSum input4.dtype=int32,int32,int32,int32 input4.format=nd,nd,nd,nd output0.name=kcacheOut output0.dtype=float16,bf16,int8,float16 output0.format=nd,nd,nd,fractal_nz output1.name=vcacheOut output1.dtype=float16,bf16,int8,float16 output1.format=nd,nd,nd,fractal_nz

该声明揭示的合法输入组合为 4 组：

1. nd / float16
2. nd / bf16
3. nd / int8
4. fractal_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/int8（910B，shape100,2,2,2，src/dst/cumSum 均为 30 长度），以及一组 310P 的block_copy_310P_fp16_nd/nz用例（如64,30,16,16的 nd 与 fractal_nz 组合），预期结果均为NO_ERROR；
• 异常用例：block_copy_ini_err（float dtype 非法，预期ERROR_INVALID_TENSOR_INI_MATCH）、block_copy_dim_diff_err（K/V 形状不一致）、block_copy_dim_err（K/V 非 4 维）、block_copy_dim_diff_err2（cumSum 与 src 形状不一致）、block_copy_index_dim_err（索引长度超过 blockCount）、block_copy_err_soc（Ascend310B 平台，预期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.csv（200+ 组随机 shape 组合）进一步扩大了 dtype（fp16/bf16/int8）、shape（blockCount 从 15 到 20000）与平台（910B/310P）的覆盖范围，其中blockCount=20000, blockSize=2, numHead=2, headSize=2这类大 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