CANN ops-transformer aclnnGatherPaKvCache 算子详解PagedAttention 非连续 KV Cache 的 Gather 拼接与两段式调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformeraclnnGatherPaKvCache 是 CANN ops-transformer 中为 PagedAttention分页注意力推理场景设计的 KV Cache 重组算子它根据 blockTables 中的 blockId 与 seqLens 中的序列长度把散落在多个物理块中、内存不连续的 token从 keyCache/valueCache 搬运并拼接成一段段连续的 key/value 序列keyRef/valueRef。本文以 aclnnGatherPaKvCache 官方接口文档 为主体骨架结合 GatherPaKvCache 算子模块 的算子定义、InferShape、Tiling 与 Kernel 源码展开讲解。读完本文你将掌握该算子的完整参数语义、isSeqLensCumsum两种 seqLens 语义下的 shape 推导规则、两段式 aclnn 接口的调用流程以及一份可直接编译运行的 C 样例。算子功能与计算逻辑功能定位内存不连续 → 内存连续在 PagedAttention 类大模型推理框架中KV Cache 以物理块block为单位按需分配。一个序列的 token 往往散布在不连续的多个物理块中无法直接作为连续的 key/value 输入喂给 Attention 内核。aclnnGatherPaKvCache 解决的就是这一gather收集环节根据blockTables中的 blockId 值、seqLens中 key/value 的 seqLen从keyCache/valueCache中将内存不连续的 token 搬运、拼接成连续的 key/value 序列。算子模块的 README 也给出了同样的定位说明见 attention/gather_pa_kv_cache/README.md。计算逻辑与输出第一维的确定规则keyRef/valueRef 的第一个维度token 总数取决于 seqLens 的内容具体由属性isSeqLensCumsum决定当isSeqLensCumsum为trueseqLens 是累加和keyRef[dim0] seqLens[-1]即 seqLens 的最后一个值就是输出序列总长度当isSeqLensCumsum为falseseqLens 是各 batch 的真实序列长度keyRef[dim0] sum(seqLens)即所有序列长度累加。该规则与 gather_pa_kv_cache_infershape.cpp 中CheckCommonPagedCacheLoad的校验逻辑相互印证is_seq_lens_cumsum为 true 时要求seqLens.shape[0] blockTables.shape[0] 1多出一个前缀 0为 false 时要求seqLens.shape[0] blockTables.shape[0]。单 token 大小限制148k 约束关于 keyRef、valueRef 有一个重要的限制条件每个 token 大小控制在148k 以内。例如对于 fp16/bf16 类型num_heads * head_sizekeyRef/valueRef取 128 × 576。这一约束与内核实现直接对应在 gather_pa_kv_cache_nd.h 中ND 内核为搬运申请了UB_BUF_SIZE 192 * 1024192KB的片上 Unified Buffer并注释ensure the size of one token less than 148kb即单 token 数据必须能在 UB 缓冲内完成中转搬运。典型 Shape 示例官方文档给出的示例 shape 如下batch16每序列最多 12 个物理块cache 块大小 12816 个 headk 的 head_size144、v 的 head_size128keyCache_shape: [128, 128, 16, 144] # [num_blocks, block_size, num_heads, head_size_k] valueCache_shape: [128, 128, 16, 128] # [num_blocks, block_size, num_heads, head_size_v] blockTables_shape: [16, 12] # [batch, block_indices] seqLens_shape: [16] # [batch] keyRef_shape: [8931, 16, 144] # [num_tokens, num_heads, head_size_k] valueRef_shape: [8931, 16, 128] # [num_tokens, num_heads, head_size_v] seqOffset_shape: [16] # [batch] out1_shape: [8931, 16, 144] out2_shape: [8931, 16, 128]其中keyRef[dim0] sum(seqLens) 8931即对应isSeqLensCumsum false的语义。产品支持情况产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册看gather_pa_kv_cache_def.cpp 中分别通过GatherPaKvCacheDefFor910b()注册ascend910b、ascend910_93两个配置和GatherPaKvCacheDefFor950()注册ascend950、ascend350两个配置完成产品适配与上表支持范围一致。两段式接口与函数原型与其他 CANN aclnn 算子一致aclnnGatherPaKvCache 采用两段式接口设计即docs/zh/context/two_phase_api.md必须先调用aclnnGatherPaKvCacheGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnGatherPaKvCache执行计算。第一段接口原型aclnnStatus aclnnGatherPaKvCacheGetWorkspaceSize( const aclTensor *keyCache, const aclTensor *valueCache, const aclTensor *blockTables, const aclTensor *seqLens, aclTensor *keyRef, aclTensor *valueRef, const aclTensor *seqOffsetOptional, char* cacheMode, bool isSeqLensCumsum, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnGatherPaKvCache( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)接口的完整实现位于 op_host/op_api/aclnn_gather_pa_kv_cache.cpp第二段接口通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成 AICore 任务下发。aclnnGatherPaKvCacheGetWorkspaceSize 参数详解下表完整列出第一段接口的全部参数来自官方文档并补充了格式细节参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorkeyCacheaclTensor*输入表示在当前层存储的 key 向量缓存当 cacheMode 为 Norm 时shape 为[num_blocks, block_size, num_heads, head_size_k]数据格式必须是 ND当 cacheMode 为 PA_NZ 时shape 为[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]数据格式必须是 FRACTAL_NZINT8、FLOAT16、BFLOAT16、FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNND、FRACTAL_NZ4√valueCacheaclTensor*输入表示在当前层存储的 value 向量缓存同 keyCache 的格式要求head_size 换成 head_size_v同 keyCacheND、FRACTAL_NZ4√blockTablesaclTensor*输入表示每个序列对应的物理块索引shape 为[batch, block_indices]其中 batch、block_indices 均须大于 0。元素取值范围为[0, num_blocks)即 blockId 可取 0 到 num_blocks-1INT32、INT64ND2×seqLensaclTensor*输入表示每个 batch 对应的序列长度shape 为[batch]或[batch 1]。当 isSeqLensCumsum 为 false 时 shape 为[batch]为 true 时 shape 为[batch 1]。元素取值范围为[0, num_blocks)与 blockTables 保持一致ND1×keyRefaclTensor*输入/输出表示 key 向量当 cacheMode 为 Norm 时 shape 为[num_tokens, num_heads, head_size_k]为 PA_NZ 时 shape 为[num_tokens, num_heads * head_size_k]与 keyCache 保持一致ND2-3√valueRefaclTensor*输入/输出表示 value 向量当 cacheMode 为 Norm 时 shape 为[num_tokens, num_heads, head_size_v]为 PA_NZ 时 shape 为[num_tokens, num_heads * head_size_v]与 valueCache 保持一致ND2-3√seqOffsetOptionalaclTensor*输入如果传入表示在从 blockTables 获取 blockId 时存在首偏移偏移量为seqOffsetOptional[i] / block_sizei表示某一个 batch不传入表示不需要偏移shape 为[batch]与 blockTables 保持一致ND1×cacheModechar*输入支持 [Norm, PA_NZ] 两种模式分别表示输入 keyCache 和 valueCache 数据格式是 ND、FRACTAL_NZ-----isSeqLensCumsumbool输入表示 seqLens 是否为累加和false 表示非累加和例如 seqLens 为[1, 3, 5, 3, 7]true 表示累加和例如 seqLens 为[0, 1, 4, 9, 12, 19]此时第 0 个元素必定是 0。累加和的seqlens[i 1] - seqlens[i]等于非累加的seqlens[i]bool---workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----属性默认值与算子定义佐证在 gather_pa_kv_cache_def.cpp 中可以看到两个属性均有默认值Attr(cache_mode).AttrType(OPTIONAL).String(Norm)cache_mode 默认取 Norm即默认按 ND 格式的 KV Cache 处理Attr(is_seq_lens_cumsum).AttrType(OPTIONAL).Bool(true)is_seq_lens_cumsum 默认取 true。Tiling 阶段gather_pa_kv_cache_tiling_arch35.cpp还会根据这两个属性 是否传入 seqOffset 组合出 8 种 tiling key1111/1110/1101/1100 对应 ND 四种组合1011/1010/1001/1000 对应 PA_NZ 四种组合见gather_pa_kv_cache_tiling.h中的tilingKeyTable实现同一算子、多内核形态的按需编译。cacheMode 对 shape 与 tiling 的影响Norm 模式NDkeyCache shape 为[num_blocks, block_size, num_heads, head_size_k]Tiling 中blockSize取自存储 shape 的第 1 维tokenSizeK num_heads * head_size_k见 gather_pa_kv_cache_tiling.cpp 的CommonGatherPaKvCacheTilingPA_NZ 模式FRACTAL_NZkeyCache shape 为[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]其中elenum_aligned与元素位宽相关——b8 场景每个数据元素位宽 8bit如 INT8取 32b16 场景如 INT16取 16b32 场景如 INT32取 8。Tiling 中blockSize取自第 2 维并校验shape[3] * 元素字节数 32B32 字节对齐见GetInputKeyCache/GetInputValueCache。平台差异与数据类型限制在支持的产品上不同架构对数据类型还有进一步限制官方文档Ascend 950PR / Ascend 950DT允许 keyCache 为 FLOAT8_E4M3FN、valueCache 为 FLOAT16 或 BFLOAT16 的组合。Atlas A2 训练/推理系列产品、Atlas A3 训练/推理系列产品输入 keyCache、valueCache、keyRef、valueRef 不支持 FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN 数据类型输入 blockTables、seqLens、seqOffsetOptional 不支持 INT64 数据类型。这与 gather_pa_kv_cache_def.cpp 的注册一致910b 系列对应 Atlas A2/A3的 key/value 仅注册 FLOAT16、BFLOAT16、INT8 三种数据类型索引类仅 INT32而 950/350 系列注册了更完整的数据类型集合。接口层aclnn_gather_pa_kv_cache.cpp还会按当前 NPU 架构DAV_3510与否选择DTYPE_SUPPORT_LIST_A5或DTYPE_SUPPORT_LIST_910B做数据校验并额外校验 keyCache 与 keyRef、valueCache 与 valueRef、blockTables 与 seqLens以及 seqOffsetOptional之间的类型一致性。返回值与错误码两段接口的返回值均为aclnnStatus具体状态码可参考 aclnn 返回码即原文档链接../../../docs/zh/context/aclnn_return_code.md。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入是空指针。ACLNN_ERR_PARAM_INVALID161002输入数据类型不在支持的范围内。ACLNN_ERR_PARAM_INVALID161002输入的维数不匹配。对应的校验逻辑在aclnnGatherPaKvCacheGetWorkspaceSize实现中依次执行CheckNullptr空指针检查→ 空 shape 提前返回 →CheckShapekeyCache/valueCache 必须 4 维、blockTables 必须 2 维、seqLens 必须 1 维、keyRef/valueRef 必须 2 或 3 维→CheckDtypeValid数据类型及一致性检查。aclnnGatherPaKvCache第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnGatherPaKvCacheGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。约束说明确定性计算aclnnGatherPaKvCache 默认确定性实现确定性计算的通用约定可参考 docs/zh/context/determinism_compute.md。源码级实现剖析Host 侧InferShape 与数据流gather_pa_kv_cache_infershape.cpp 负责 shape 推导输出 key/value 的 shape 直接继承输入 key/value第 4、5 个输入的 shape输入存在 UnknownRank/UnknownShape动态场景时输出相应置为未知校验 blockTables 维数为 2、seqLens 维数为 1并按is_seq_lens_cumsum校验seqLens.shape[0]与blockTables.shape[0]的关系truebatch1falsebatch按 cache_mode 分别走InferShape4GatherPaKvCacheNdkey/value 输出为 3 维或InferShape4GatherPaKvCacheNzkey/value 输出为 2 维。Host 侧Tiling 与分核gather_pa_kv_cache_tiling.cpp 与 gather_pa_kv_cache_tiling_arch35.cpp 完成计算切分通用路径固定申请 16MB workspaceASCENDC_TOOLS_WORKSPACEblockDim分核数Norm 模式使用全部 AIV 核数PA_NZ 模式取min(总token数, AIV核数)arch35 路径针对非连续 tensorview场景做了专门设计通过GetTensorInfo获取逻辑 shape 与 stride判断 cache/ref 各轴是否连续用nonContiguousFlag位图标记非连续形态bit0~bit9 分别对应 keyCache、valueCache、key 输出、value 输出及各轴的非连续状态并将各 stride 写入 tiling 数据ND 非连续场景改为全局 block 轮询分核解决 batch 数很少时按 batch 分核只用少量核的负载均衡问题。Kernel 侧AIV 搬运内核gather_pa_kv_cache.cpp 是内核入口声明为KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)按编译期 dtype 与 tiling key 分发到ND 内核GatherPaKvCacheNdTgather_pa_kv_cache_nd.htiling key 618INT8/619FLOAT16、BF16NZ 内核GatherPaKvCacheNzTgather_pa_kv_cache_nz.htiling key 577。ND 内核中GetBlockCacheOffset根据cacheBlockStrideview 场景下由 host 传入的 stride计算物理块首地址偏移随后按blockSize × tokenSize分 cacheline 切块搬运到 192KB UB 缓冲再写回连续输出这正是非连续 token 搬运拼接的落地实现。接口层非连续 Tensor 的零拷贝与兜底aclnn_gather_pa_kv_cache.cpp 在入参校验后会根据各 tensor 的 view shape/stride 判断非连续形态并分流ProcessNonContiguouscache 首轴非连续dim0或 ND 模式内部轴slot/head非连续、ref 内部轴非连续但尾轴连续等 kernel 可零拷贝处理的形态通过CreateView保留 view stride 直接下发给 kernelkernel 按 stride 搬运/散写避免物理拷贝ProcessContiguous对 kernel 无法处理的形态如 NZ 非 dim0 轴非连续、ND 尾轴非连续先通过l0op::Contiguous物理连续化计算后经ViewCopy回写到原始 ref 内存保证语义正确。这也解释了文档参数表中 keyCache/valueCache/keyRef/valueRef 均标注非连续 Tensor √——该算子对非连续输入/输出有完整支持。调用示例可直接编译运行的 C 样例下面示例取自官方文档与仓库 examples/test_aclnn_gather_pa_kv_cache.cpp 等价演示了从环境初始化、构造输入输出、两段式调用到结果回拷与资源释放的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_gather_pa_kv_cache.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. (固定写法)device/stream初始化参考acl对外接口列表 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t keyCacheShape {2, 2, 32, 2}; std::vectorint64_t valueCacheShape {2, 2, 32, 4}; std::vectorint64_t blockTablesShape {4,6}; std::vectorint64_t seqLensShape {4}; std::vectorint64_t keyShape {12, 32, 2}; std::vectorint64_t valueShape {12, 32, 4}; std::vectorint64_t seqOffsetShape {4}; void* keyCacheDeviceAddr nullptr; void* valueCacheDeviceAddr nullptr; void* blockTablesDeviceAddr nullptr; void* seqLensDeviceAddr nullptr; void* keyDeviceAddr nullptr; void* valueDeviceAddr nullptr; void* seqOffsetAddr nullptr; aclTensor* keyCache nullptr; aclTensor* valueCache nullptr; aclTensor* blockTables nullptr; aclTensor* seqLens nullptr; aclTensor* key nullptr; aclTensor* value nullptr; aclTensor* seqOffset nullptr; std::vectoruint16_t keyCacheHostData(256, 1); std::vectoruint16_t valueCacheHostData(512, 1); std::vectorint32_t blockTablesHostData(24, 1); std::vectorint32_t seqLensHostData(4, 3); std::vectoruint16_t keyHostData(768, 0); std::vectoruint16_t valueHostData(1536, 0); std::vectorint32_t seqOffsetHostData(4, 2); char cacheMode[] Norm; const bool isSeqLensCumsum false; // 创建GatherPaKvCache的输入输出aclTensor ret CreateAclTensor(keyCacheHostData, keyCacheShape, keyCacheDeviceAddr, aclDataType::ACL_FLOAT16, keyCache); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(valueCacheHostData, valueCacheShape, valueCacheDeviceAddr, aclDataType::ACL_FLOAT16, valueCache); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(blockTablesHostData, blockTablesShape, blockTablesDeviceAddr, aclDataType::ACL_INT32, blockTables); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(seqLensHostData, seqLensShape, seqLensDeviceAddr, aclDataType::ACL_INT32, seqLens); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(keyHostData, keyShape, keyDeviceAddr, aclDataType::ACL_FLOAT16, key); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(valueHostData, valueShape, valueDeviceAddr, aclDataType::ACL_FLOAT16, value); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(seqOffsetHostData, seqOffsetShape, seqOffsetAddr, aclDataType::ACL_INT32, seqOffset); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的Api名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnGatherPaKvCache第一段接口 ret aclnnGatherPaKvCacheGetWorkspaceSize(keyCache, valueCache, blockTables, seqLens, key , value, seqOffset, cacheMode, isSeqLensCumsum, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnGatherPaKvCacheGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnGatherPaKvCache第二段接口 ret aclnnGatherPaKvCache(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnGatherPaKvCache failed. ERROR: %d\n, ret); return ret); // 4. (固定写法)同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(keyShape); std::vectoruint16_t resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), keyDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %d\n, i, resultData[i]); } // 6. 释放aclTensor和aclIntArray需要根据具体API的接口定义修改 aclDestroyTensor(keyCache); aclDestroyTensor(valueCache); aclDestroyTensor(blockTables); aclDestroyTensor(seqLens); aclDestroyTensor(key); aclDestroyTensor(value); aclDestroyTensor(seqOffset); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(keyCacheDeviceAddr); aclrtFree(valueCacheDeviceAddr); aclrtFree(blockTablesDeviceAddr ); aclrtFree(seqLensDeviceAddr ); aclrtFree(keyDeviceAddr); aclrtFree(valueDeviceAddr); aclrtFree(seqOffsetAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点解读示例中keyCacheShape {2, 2, 32, 2}对应 Norm 模式的[num_blocks2, block_size2, num_heads32, head_size_k2]seqLensHostData {3, 3, 3, 3}非累加和isSeqLensCumsum false因此输出keyShape {12, 32, 2}12 4 * 3 sum(seqLens)seqOffsetHostData {2, 2, 2, 2}表示每个 batch 从blockTables取 blockId 时存在首偏移偏移块数 seqOffset[i] / block_size 2 / 2 1个块workspace 按第一段接口返回的大小可能为 0条件性申请完整示例还可参考仓库 examples/test_aclnn_gather_pa_kv_cache.cpp以及算子模块下的单元测试 tests/ut覆盖 InferShape、Tiling 与 Kernel 三层。编译与运行该示例属于标准 CANN aclnn 单算子调用工程具体编译、链接与运行方式包括头文件路径、libascendcl/libopapi等库的链接以及 NPU 环境变量配置请参考仓库通用文档 编译与运行样例即原文档链接../../../docs/zh/context/compile_and_run_sample.md。运行前需确保宿主机已安装与目标产品匹配的 CANN Toolkit 与算子包且设备为上文产品支持情况中列出的支持型号编译时链接 aclnn 相关动态库并正确包含aclnnop/aclnn_gather_pa_kv_cache.h头文件根据实际设备填写deviceId并根据实际的 KV Cache 规模调整各 tensor 的 shape 与数据类型。小结aclnnGatherPaKvCache 是 CANN ops-transformer 面向 PagedAttention 推理的关键算子之一。本文围绕 官方接口文档 完整梳理了其功能语义blockTables seqLens 驱动的非连续 KV Cache 搬运拼接、isSeqLensCumsum两种 seqLens 语义、两段式接口的全部参数与错误码、平台差异与数据类型限制并深入 算子模块源码 印证了 InferShape、Tiling 分核、ND/NZ 内核搬运及非连续 tensor 零拷贝处理等底层实现最后给出了一份可直接落地的完整 C 调用示例可作为在 NPU 上接入该算子的直接参考。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-transformer aclnnGatherPaKvCache 算子详解：PagedAttention 非连续 KV Cache 的 Gather 拼接与两段式调用实战
CANN ops-transformer aclnnGatherPaKvCache 算子详解PagedAttention 非连续 KV Cache 的 Gather 拼接与两段式调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformeraclnnGatherPaKvCache 是 CANN ops-transformer 中为 PagedAttention分页注意力推理场景设计的 KV Cache 重组算子它根据 blockTables 中的 blockId 与 seqLens 中的序列长度把散落在多个物理块中、内存不连续的 token从 keyCache/valueCache 搬运并拼接成一段段连续的 key/value 序列keyRef/valueRef。本文以 aclnnGatherPaKvCache 官方接口文档 为主体骨架结合 GatherPaKvCache 算子模块 的算子定义、InferShape、Tiling 与 Kernel 源码展开讲解。读完本文你将掌握该算子的完整参数语义、isSeqLensCumsum两种 seqLens 语义下的 shape 推导规则、两段式 aclnn 接口的调用流程以及一份可直接编译运行的 C 样例。算子功能与计算逻辑功能定位内存不连续 → 内存连续在 PagedAttention 类大模型推理框架中KV Cache 以物理块block为单位按需分配。一个序列的 token 往往散布在不连续的多个物理块中无法直接作为连续的 key/value 输入喂给 Attention 内核。aclnnGatherPaKvCache 解决的就是这一gather收集环节根据blockTables中的 blockId 值、seqLens中 key/value 的 seqLen从keyCache/valueCache中将内存不连续的 token 搬运、拼接成连续的 key/value 序列。算子模块的 README 也给出了同样的定位说明见 attention/gather_pa_kv_cache/README.md。计算逻辑与输出第一维的确定规则keyRef/valueRef 的第一个维度token 总数取决于 seqLens 的内容具体由属性isSeqLensCumsum决定当isSeqLensCumsum为trueseqLens 是累加和keyRef[dim0] seqLens[-1]即 seqLens 的最后一个值就是输出序列总长度当isSeqLensCumsum为falseseqLens 是各 batch 的真实序列长度keyRef[dim0] sum(seqLens)即所有序列长度累加。该规则与 gather_pa_kv_cache_infershape.cpp 中CheckCommonPagedCacheLoad的校验逻辑相互印证is_seq_lens_cumsum为 true 时要求seqLens.shape[0] blockTables.shape[0] 1多出一个前缀 0为 false 时要求seqLens.shape[0] blockTables.shape[0]。单 token 大小限制148k 约束关于 keyRef、valueRef 有一个重要的限制条件每个 token 大小控制在148k 以内。例如对于 fp16/bf16 类型num_heads * head_sizekeyRef/valueRef取 128 × 576。这一约束与内核实现直接对应在 gather_pa_kv_cache_nd.h 中ND 内核为搬运申请了UB_BUF_SIZE 192 * 1024192KB的片上 Unified Buffer并注释ensure the size of one token less than 148kb即单 token 数据必须能在 UB 缓冲内完成中转搬运。典型 Shape 示例官方文档给出的示例 shape 如下batch16每序列最多 12 个物理块cache 块大小 12816 个 headk 的 head_size144、v 的 head_size128keyCache_shape: [128, 128, 16, 144] # [num_blocks, block_size, num_heads, head_size_k] valueCache_shape: [128, 128, 16, 128] # [num_blocks, block_size, num_heads, head_size_v] blockTables_shape: [16, 12] # [batch, block_indices] seqLens_shape: [16] # [batch] keyRef_shape: [8931, 16, 144] # [num_tokens, num_heads, head_size_k] valueRef_shape: [8931, 16, 128] # [num_tokens, num_heads, head_size_v] seqOffset_shape: [16] # [batch] out1_shape: [8931, 16, 144] out2_shape: [8931, 16, 128]其中keyRef[dim0] sum(seqLens) 8931即对应isSeqLensCumsum false的语义。产品支持情况产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从算子注册看gather_pa_kv_cache_def.cpp 中分别通过GatherPaKvCacheDefFor910b()注册ascend910b、ascend910_93两个配置和GatherPaKvCacheDefFor950()注册ascend950、ascend350两个配置完成产品适配与上表支持范围一致。两段式接口与函数原型与其他 CANN aclnn 算子一致aclnnGatherPaKvCache 采用两段式接口设计即docs/zh/context/two_phase_api.md必须先调用aclnnGatherPaKvCacheGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnGatherPaKvCache执行计算。第一段接口原型aclnnStatus aclnnGatherPaKvCacheGetWorkspaceSize( const aclTensor *keyCache, const aclTensor *valueCache, const aclTensor *blockTables, const aclTensor *seqLens, aclTensor *keyRef, aclTensor *valueRef, const aclTensor *seqOffsetOptional, char* cacheMode, bool isSeqLensCumsum, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnGatherPaKvCache( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)接口的完整实现位于 op_host/op_api/aclnn_gather_pa_kv_cache.cpp第二段接口通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成 AICore 任务下发。aclnnGatherPaKvCacheGetWorkspaceSize 参数详解下表完整列出第一段接口的全部参数来自官方文档并补充了格式细节参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorkeyCacheaclTensor*输入表示在当前层存储的 key 向量缓存当 cacheMode 为 Norm 时shape 为[num_blocks, block_size, num_heads, head_size_k]数据格式必须是 ND当 cacheMode 为 PA_NZ 时shape 为[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]数据格式必须是 FRACTAL_NZINT8、FLOAT16、BFLOAT16、FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FNND、FRACTAL_NZ4√valueCacheaclTensor*输入表示在当前层存储的 value 向量缓存同 keyCache 的格式要求head_size 换成 head_size_v同 keyCacheND、FRACTAL_NZ4√blockTablesaclTensor*输入表示每个序列对应的物理块索引shape 为[batch, block_indices]其中 batch、block_indices 均须大于 0。元素取值范围为[0, num_blocks)即 blockId 可取 0 到 num_blocks-1INT32、INT64ND2×seqLensaclTensor*输入表示每个 batch 对应的序列长度shape 为[batch]或[batch 1]。当 isSeqLensCumsum 为 false 时 shape 为[batch]为 true 时 shape 为[batch 1]。元素取值范围为[0, num_blocks)与 blockTables 保持一致ND1×keyRefaclTensor*输入/输出表示 key 向量当 cacheMode 为 Norm 时 shape 为[num_tokens, num_heads, head_size_k]为 PA_NZ 时 shape 为[num_tokens, num_heads * head_size_k]与 keyCache 保持一致ND2-3√valueRefaclTensor*输入/输出表示 value 向量当 cacheMode 为 Norm 时 shape 为[num_tokens, num_heads, head_size_v]为 PA_NZ 时 shape 为[num_tokens, num_heads * head_size_v]与 valueCache 保持一致ND2-3√seqOffsetOptionalaclTensor*输入如果传入表示在从 blockTables 获取 blockId 时存在首偏移偏移量为seqOffsetOptional[i] / block_sizei表示某一个 batch不传入表示不需要偏移shape 为[batch]与 blockTables 保持一致ND1×cacheModechar*输入支持 [Norm, PA_NZ] 两种模式分别表示输入 keyCache 和 valueCache 数据格式是 ND、FRACTAL_NZ-----isSeqLensCumsumbool输入表示 seqLens 是否为累加和false 表示非累加和例如 seqLens 为[1, 3, 5, 3, 7]true 表示累加和例如 seqLens 为[0, 1, 4, 9, 12, 19]此时第 0 个元素必定是 0。累加和的seqlens[i 1] - seqlens[i]等于非累加的seqlens[i]bool---workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----属性默认值与算子定义佐证在 gather_pa_kv_cache_def.cpp 中可以看到两个属性均有默认值Attr(cache_mode).AttrType(OPTIONAL).String(Norm)cache_mode 默认取 Norm即默认按 ND 格式的 KV Cache 处理Attr(is_seq_lens_cumsum).AttrType(OPTIONAL).Bool(true)is_seq_lens_cumsum 默认取 true。Tiling 阶段gather_pa_kv_cache_tiling_arch35.cpp还会根据这两个属性 是否传入 seqOffset 组合出 8 种 tiling key1111/1110/1101/1100 对应 ND 四种组合1011/1010/1001/1000 对应 PA_NZ 四种组合见gather_pa_kv_cache_tiling.h中的tilingKeyTable实现同一算子、多内核形态的按需编译。cacheMode 对 shape 与 tiling 的影响Norm 模式NDkeyCache shape 为[num_blocks, block_size, num_heads, head_size_k]Tiling 中blockSize取自存储 shape 的第 1 维tokenSizeK num_heads * head_size_k见 gather_pa_kv_cache_tiling.cpp 的CommonGatherPaKvCacheTilingPA_NZ 模式FRACTAL_NZkeyCache shape 为[num_blocks, num_heads * head_size_k // elenum_aligned, block_size, elenum_aligned]其中elenum_aligned与元素位宽相关——b8 场景每个数据元素位宽 8bit如 INT8取 32b16 场景如 INT16取 16b32 场景如 INT32取 8。Tiling 中blockSize取自第 2 维并校验shape[3] * 元素字节数 32B32 字节对齐见GetInputKeyCache/GetInputValueCache。平台差异与数据类型限制在支持的产品上不同架构对数据类型还有进一步限制官方文档Ascend 950PR / Ascend 950DT允许 keyCache 为 FLOAT8_E4M3FN、valueCache 为 FLOAT16 或 BFLOAT16 的组合。Atlas A2 训练/推理系列产品、Atlas A3 训练/推理系列产品输入 keyCache、valueCache、keyRef、valueRef 不支持 FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN 数据类型输入 blockTables、seqLens、seqOffsetOptional 不支持 INT64 数据类型。这与 gather_pa_kv_cache_def.cpp 的注册一致910b 系列对应 Atlas A2/A3的 key/value 仅注册 FLOAT16、BFLOAT16、INT8 三种数据类型索引类仅 INT32而 950/350 系列注册了更完整的数据类型集合。接口层aclnn_gather_pa_kv_cache.cpp还会按当前 NPU 架构DAV_3510与否选择DTYPE_SUPPORT_LIST_A5或DTYPE_SUPPORT_LIST_910B做数据校验并额外校验 keyCache 与 keyRef、valueCache 与 valueRef、blockTables 与 seqLens以及 seqOffsetOptional之间的类型一致性。返回值与错误码两段接口的返回值均为aclnnStatus具体状态码可参考 aclnn 返回码即原文档链接../../../docs/zh/context/aclnn_return_code.md。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入是空指针。ACLNN_ERR_PARAM_INVALID161002输入数据类型不在支持的范围内。ACLNN_ERR_PARAM_INVALID161002输入的维数不匹配。对应的校验逻辑在aclnnGatherPaKvCacheGetWorkspaceSize实现中依次执行CheckNullptr空指针检查→ 空 shape 提前返回 →CheckShapekeyCache/valueCache 必须 4 维、blockTables 必须 2 维、seqLens 必须 1 维、keyRef/valueRef 必须 2 或 3 维→CheckDtypeValid数据类型及一致性检查。aclnnGatherPaKvCache第二段接口参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnGatherPaKvCacheGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。约束说明确定性计算aclnnGatherPaKvCache 默认确定性实现确定性计算的通用约定可参考 docs/zh/context/determinism_compute.md。源码级实现剖析Host 侧InferShape 与数据流gather_pa_kv_cache_infershape.cpp 负责 shape 推导输出 key/value 的 shape 直接继承输入 key/value第 4、5 个输入的 shape输入存在 UnknownRank/UnknownShape动态场景时输出相应置为未知校验 blockTables 维数为 2、seqLens 维数为 1并按is_seq_lens_cumsum校验seqLens.shape[0]与blockTables.shape[0]的关系truebatch1falsebatch按 cache_mode 分别走InferShape4GatherPaKvCacheNdkey/value 输出为 3 维或InferShape4GatherPaKvCacheNzkey/value 输出为 2 维。Host 侧Tiling 与分核gather_pa_kv_cache_tiling.cpp 与 gather_pa_kv_cache_tiling_arch35.cpp 完成计算切分通用路径固定申请 16MB workspaceASCENDC_TOOLS_WORKSPACEblockDim分核数Norm 模式使用全部 AIV 核数PA_NZ 模式取min(总token数, AIV核数)arch35 路径针对非连续 tensorview场景做了专门设计通过GetTensorInfo获取逻辑 shape 与 stride判断 cache/ref 各轴是否连续用nonContiguousFlag位图标记非连续形态bit0~bit9 分别对应 keyCache、valueCache、key 输出、value 输出及各轴的非连续状态并将各 stride 写入 tiling 数据ND 非连续场景改为全局 block 轮询分核解决 batch 数很少时按 batch 分核只用少量核的负载均衡问题。Kernel 侧AIV 搬运内核gather_pa_kv_cache.cpp 是内核入口声明为KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)按编译期 dtype 与 tiling key 分发到ND 内核GatherPaKvCacheNdTgather_pa_kv_cache_nd.htiling key 618INT8/619FLOAT16、BF16NZ 内核GatherPaKvCacheNzTgather_pa_kv_cache_nz.htiling key 577。ND 内核中GetBlockCacheOffset根据cacheBlockStrideview 场景下由 host 传入的 stride计算物理块首地址偏移随后按blockSize × tokenSize分 cacheline 切块搬运到 192KB UB 缓冲再写回连续输出这正是非连续 token 搬运拼接的落地实现。接口层非连续 Tensor 的零拷贝与兜底aclnn_gather_pa_kv_cache.cpp 在入参校验后会根据各 tensor 的 view shape/stride 判断非连续形态并分流ProcessNonContiguouscache 首轴非连续dim0或 ND 模式内部轴slot/head非连续、ref 内部轴非连续但尾轴连续等 kernel 可零拷贝处理的形态通过CreateView保留 view stride 直接下发给 kernelkernel 按 stride 搬运/散写避免物理拷贝ProcessContiguous对 kernel 无法处理的形态如 NZ 非 dim0 轴非连续、ND 尾轴非连续先通过l0op::Contiguous物理连续化计算后经ViewCopy回写到原始 ref 内存保证语义正确。这也解释了文档参数表中 keyCache/valueCache/keyRef/valueRef 均标注非连续 Tensor √——该算子对非连续输入/输出有完整支持。调用示例可直接编译运行的 C 样例下面示例取自官方文档与仓库 examples/test_aclnn_gather_pa_kv_cache.cpp 等价演示了从环境初始化、构造输入输出、两段式调用到结果回拷与资源释放的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_gather_pa_kv_cache.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. (固定写法)device/stream初始化参考acl对外接口列表 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t keyCacheShape {2, 2, 32, 2}; std::vectorint64_t valueCacheShape {2, 2, 32, 4}; std::vectorint64_t blockTablesShape {4,6}; std::vectorint64_t seqLensShape {4}; std::vectorint64_t keyShape {12, 32, 2}; std::vectorint64_t valueShape {12, 32, 4}; std::vectorint64_t seqOffsetShape {4}; void* keyCacheDeviceAddr nullptr; void* valueCacheDeviceAddr nullptr; void* blockTablesDeviceAddr nullptr; void* seqLensDeviceAddr nullptr; void* keyDeviceAddr nullptr; void* valueDeviceAddr nullptr; void* seqOffsetAddr nullptr; aclTensor* keyCache nullptr; aclTensor* valueCache nullptr; aclTensor* blockTables nullptr; aclTensor* seqLens nullptr; aclTensor* key nullptr; aclTensor* value nullptr; aclTensor* seqOffset nullptr; std::vectoruint16_t keyCacheHostData(256, 1); std::vectoruint16_t valueCacheHostData(512, 1); std::vectorint32_t blockTablesHostData(24, 1); std::vectorint32_t seqLensHostData(4, 3); std::vectoruint16_t keyHostData(768, 0); std::vectoruint16_t valueHostData(1536, 0); std::vectorint32_t seqOffsetHostData(4, 2); char cacheMode[] Norm; const bool isSeqLensCumsum false; // 创建GatherPaKvCache的输入输出aclTensor ret CreateAclTensor(keyCacheHostData, keyCacheShape, keyCacheDeviceAddr, aclDataType::ACL_FLOAT16, keyCache); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(valueCacheHostData, valueCacheShape, valueCacheDeviceAddr, aclDataType::ACL_FLOAT16, valueCache); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(blockTablesHostData, blockTablesShape, blockTablesDeviceAddr, aclDataType::ACL_INT32, blockTables); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(seqLensHostData, seqLensShape, seqLensDeviceAddr, aclDataType::ACL_INT32, seqLens); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(keyHostData, keyShape, keyDeviceAddr, aclDataType::ACL_FLOAT16, key); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(valueHostData, valueShape, valueDeviceAddr, aclDataType::ACL_FLOAT16, value); CHECK_RET(ret ACL_SUCCESS, return ret); ret CreateAclTensor(seqOffsetHostData, seqOffsetShape, seqOffsetAddr, aclDataType::ACL_INT32, seqOffset); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的Api名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnGatherPaKvCache第一段接口 ret aclnnGatherPaKvCacheGetWorkspaceSize(keyCache, valueCache, blockTables, seqLens, key , value, seqOffset, cacheMode, isSeqLensCumsum, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnGatherPaKvCacheGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnGatherPaKvCache第二段接口 ret aclnnGatherPaKvCache(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnGatherPaKvCache failed. ERROR: %d\n, ret); return ret); // 4. (固定写法)同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(keyShape); std::vectoruint16_t resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), keyDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %d\n, i, resultData[i]); } // 6. 释放aclTensor和aclIntArray需要根据具体API的接口定义修改 aclDestroyTensor(keyCache); aclDestroyTensor(valueCache); aclDestroyTensor(blockTables); aclDestroyTensor(seqLens); aclDestroyTensor(key); aclDestroyTensor(value); aclDestroyTensor(seqOffset); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(keyCacheDeviceAddr); aclrtFree(valueCacheDeviceAddr); aclrtFree(blockTablesDeviceAddr ); aclrtFree(seqLensDeviceAddr ); aclrtFree(keyDeviceAddr); aclrtFree(valueDeviceAddr); aclrtFree(seqOffsetAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点解读示例中keyCacheShape {2, 2, 32, 2}对应 Norm 模式的[num_blocks2, block_size2, num_heads32, head_size_k2]seqLensHostData {3, 3, 3, 3}非累加和isSeqLensCumsum false因此输出keyShape {12, 32, 2}12 4 * 3 sum(seqLens)seqOffsetHostData {2, 2, 2, 2}表示每个 batch 从blockTables取 blockId 时存在首偏移偏移块数 seqOffset[i] / block_size 2 / 2 1个块workspace 按第一段接口返回的大小可能为 0条件性申请完整示例还可参考仓库 examples/test_aclnn_gather_pa_kv_cache.cpp以及算子模块下的单元测试 tests/ut覆盖 InferShape、Tiling 与 Kernel 三层。编译与运行该示例属于标准 CANN aclnn 单算子调用工程具体编译、链接与运行方式包括头文件路径、libascendcl/libopapi等库的链接以及 NPU 环境变量配置请参考仓库通用文档 编译与运行样例即原文档链接../../../docs/zh/context/compile_and_run_sample.md。运行前需确保宿主机已安装与目标产品匹配的 CANN Toolkit 与算子包且设备为上文产品支持情况中列出的支持型号编译时链接 aclnn 相关动态库并正确包含aclnnop/aclnn_gather_pa_kv_cache.h头文件根据实际设备填写deviceId并根据实际的 KV Cache 规模调整各 tensor 的 shape 与数据类型。小结aclnnGatherPaKvCache 是 CANN ops-transformer 面向 PagedAttention 推理的关键算子之一。本文围绕 官方接口文档 完整梳理了其功能语义blockTables seqLens 驱动的非连续 KV Cache 搬运拼接、isSeqLensCumsum两种 seqLens 语义、两段式接口的全部参数与错误码、平台差异与数据类型限制并深入 算子模块源码 印证了 InferShape、Tiling 分核、ND/NZ 内核搬运及非连续 tensor 零拷贝处理等底层实现最后给出了一份可直接落地的完整 C 调用示例可作为在 NPU 上接入该算子的直接参考。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考