CANN ops-transformer 算子解析aclnnScatterPaKvCache 接口详解与调用实战【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformeraclnnScatterPaKvCache 是 CANN ops-transformer 仓库中用于 PagedAttention 推理场景的关键算子接口其功能是按 slotMapping 指定的偏移将当前 step 多个 token 的 key/value 散射Scatter写入 KV Cache 的指定位置。本文以 docs/aclnnScatterPaKvCache.md 为主体结合仓库中 op_host、op_kernel 的源码完整讲解该算子的 7 种输入输出场景、两段式接口调用方式、全部参数与约束以及从 host 侧校验到 kernel 分发的底层执行链路帮助读者在 Atlas A2/A3 与 Ascend 950 系列产品上正确、高效地完成 KV Cache 写入。功能说明与适用场景该算子的接口功能为更新 KvCache 中指定位置的 key 和 value。它接收当前 step 待写入的 key/value 张量与整块 KV CachekeyCacheRef / valueCacheRef通过 slotMapping 中给出的token 在 cache 中的存储偏移把每个 token 的 key/value 写入 cache 对应位置。该过程属于典型的 gather/scatter 式 PagedAttention 前处理常见于解码阶段每一轮生成后的 KV 落盘环节。从算子属性角度看其行为由两个关键属性决定定义见 scatter_pa_kv_cache_def.cppcache_mode默认Norm决定 keyCacheRef / valueCacheRef 的内存排布格式可选等价 Norm、NormND 排布、PA_NZFRACTAL_NZ 排布。scatter_mode默认None决定待写入的 key/value 的状态可选等价 None、None、Alibi、Rope、Omni、Nct、NHSD。这两个属性在 host 侧的 tiling 阶段会被映射为具体的内核模板类型在 scatter_pa_kv_cache_tiling.cpp 的GetTemplateType()中完成分发PA_NZ走 NZ 模板Norm下分别按Rope/Alibi/Omni/Nct/NHSD/None选择对应模板。产品支持情况产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持注意产品间存在场景差异Atlas A3 / A2 系列仅支持场景一、二、四、五、六、七即不支持带 4 维 key/value 的压缩场景三。Ascend 950PR / 950DT仅支持场景一、二、三、四、五即不支持 SISO 场景六与 NHSD 场景七。同时不同产品对 key/value/keyCacheRef/valueCacheRef 的数据类型支持范围也不同详见下文参数详解。输入输出支持的七种场景算子通过入参构造自动区分场景符合第一种入参构造走场景一符合第二种构造走场景二依次类推。场景一、场景二、场景六没有compressLensOptional、seqLensOptional、compressSeqOffsetOptional这三个可选参数场景四没有compressSeqOffsetOptional可选参数。场景一PA_NZ 排布的 cachekey:[batch * seq_len, num_head, k_head_size] value:[batch * seq_len, num_head, v_head_size] keyCache:[num_blocks, num_head * k_head_size // last_dim_k, block_size, last_dim_k]/[num_blocks, num_head, k_head_size // last_dim_k, block_size, last_dim_k] valueCache:[num_blocks, num_head * v_head_size // last_dim_v, block_size, last_dim_v]/[num_blocks, num_head, v_head_size // last_dim_v, block_size, last_dim_v] slotMapping:[batch * seq_len] cacheMode:PA_NZ last_dim_k 32 / sizeof(dtypeKey) last_dim_v 32 / sizeof(dtypeValue) (k_head_size * sizeof(dtypeKey)) % 32 0 (v_head_size * sizeof(dtypeValue)) % 32 0该场景要求 cache 采用 FRACTAL_NZ 内存排布keyCacheRef/valueCacheRef 既支持 4 维[num_blocks, num_head*k_head_size//last_dim_k, block_size, last_dim_k]这种把 num_head 与 head 维合并的紧凑排布也支持 5 维[num_blocks, num_head, k_head_size//last_dim_k, block_size, last_dim_k]。tiling 侧对应TEMPLATE_NZ值为 2并额外校验blockSize UINT16_MAX、key/value 的 head_size 需 32 字节对齐见 scatter_pa_kv_cache_tiling.cpp 的CheckInputShapeNz()。场景二Norm 排布的常规写入key:[batch * seq_len, num_head, k_head_size] value:[batch * seq_len, num_head, v_head_size] keyCache:[num_blocks, block_size, num_head, k_head_size] valueCache:[num_blocks, block_size, num_head, v_head_size] slotMapping:[batch * seq_len] cacheMode:Norm scatter_mode:None/Nct这是最常用的 ND 排布场景其中 k_head_size 与 v_head_size 可以不同也可以相同。scatter_mode为None时表示 key/value 是连续张量为Nct时表示 key/value 非连续此时需配合 strides/offsets 参数见下文。tiling 侧分别对应TEMPLATE_NORMAL值 1与TEMPLATE_NORM_NCT值 7。场景三4 维 key/value 的分 batch、分 head 压缩写入仅 950 系列key:[batch, seq_len, num_head, k_head_size] value:[batch, seq_len, num_head, v_head_size] keyCache:[num_blocks, block_size, 1, k_head_size] valueCache:[num_blocks, block_size, 1, v_head_size] slotMapping:[batch, num_head] compressLensOptional:[batch, num_head] seqLensOptional:[batch] compressSeqOffsetOptional:[batch * num_head] cacheMode:Norm该场景 key/value 为 4 维slotMapping 为 2 维[batch, num_head]cache 的第三维固定为 1表示 cache 按 batch 顺序铺排不分 head 独立分块配合compressLensOptional压缩量与seqLensOptional每个 batch 的实际序列长度实现带压缩的写入。此场景仅有 Ascend 950 系列支持。场景四Alibi 压缩结构写入key:[num_tokens, num_head, k_head_size] value:[num_tokens, num_head, v_head_size] keyCache:[num_blocks, block_size, 1, k_head_size] valueCache:[num_blocks, block_size, 1, v_head_size] slotMapping:[batch * num_head] compressLensOptional:[batch * num_head] seqLensOptional:[batch] cacheMode:Norm scatter_mode:Alibikey/value 为 3 维num_tokens为所有 batch 的 token 总数scatter_mode为Alibi时表示更新的是基于 Alibi 结构的压缩状态 key/value。该场景没有compressSeqOffsetOptional参数。tiling 对应TEMPLATE_ALIBI值 3其分核策略为按 token 数切分见 scatter_pa_kv_cache_tiling.cpp。场景五Rope / Omni 压缩结构写入key:[num_tokens, num_head, k_head_size] value:[num_tokens, num_head, v_head_size] keyCache:[num_blocks, block_size, 1, k_head_size] valueCache:[num_blocks, block_size, 1, v_head_size] slotMapping:[batch * num_head] compressLensOptional:[batch * num_head] seqLensOptional:[batch] compressSeqOffsetOptional:[batch * num_head] cacheMode:Norm scatter_mode:Rope/Omni与场景四类似但scatter_mode为Rope或Omni且带compressSeqOffsetOptional每个 batch 每个 head 的压缩起点。tiling 对应TEMPLATE_ROPE值 4与TEMPLATE_OMNI值 6。这两个模板在分核时使用TASK_MULTIPLE 2即 KV 分核任务量翻倍见 scatter_pa_kv_cache_tiling.cpp且仅支持 FLOAT16 / BFLOAT16 数据类型kernel 侧通过宏ORIG_DTYPE_KEY DT_FLOAT16 || DT_BF16控制编译见 scatter_pa_kv_cache.cpp。场景六仅写 key 的 SISO 场景key:[batch * seq_len, num_head, k_head_size] value:[] keyCache:[num_blocks, block_size, num_head, k_head_size] valueCache:[] slotMapping:[batch * seq_len] cacheMode:Norm scatter_mode:None/Nctvalue 与 valueCacheRef 均为 0 维空张量表示只更新 key 不更新 valueSingle-Input-Single-OutputSISO。tiling 侧在GetTemplateType()中检测到 valueCache 维度为 0 时将TEMPLATE_NORMAL改写为TEMPLATE_SISO值 5、TEMPLATE_NORM_NCT改写为TEMPLATE_SISO_NCT值 8kernel 对应ScatterPaKvCacheNormalSiso实现见 scatter_pa_kv_cache.cpp。场景七NHSD 排布写入key:[num_tokens, num_head, k_head_size] value:[num_tokens, num_head, v_head_size] keyCache:[num_blocks, num_head, block_size, k_head_size] valueCache:[num_blocks, num_head, block_size, v_head_size] slotMapping:[num_tokens] cacheMode:Norm scatter_mode:NHSDcache 采用[num_blocks, num_head, block_size, head_size]NHSD即 num_head 位于第二维的排布slotMapping 为一维[num_tokens]。tiling 对应TEMPLATE_NHSD值 9并有一组专属约束详见下文约束说明。函数原型与两段式接口与 CANN 其他 aclnn 算子一致aclnnScatterPaKvCache 采用两段式接口详见 两段式接口说明必须先调用aclnnScatterPaKvCacheGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnScatterPaKvCache执行计算。aclnnStatus aclnnScatterPaKvCacheGetWorkspaceSize( const aclTensor *key, aclTensor *keyCacheRef, const aclTensor *slotMapping, const aclTensor *value, aclTensor *valueCacheRef, const aclTensor *compressLensOptional, const aclTensor *compressSeqOffsetOptional, const aclTensor *seqLensOptional, char *cacheModeOptional, char *scatterModeOptional, const aclIntArray *stridesOptional, const aclIntArray *offsetsOptional, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnScatterPaKvCache( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从仓库实现看第一段接口除了常规的空指针与数据类型/形状校验外还做了两类特殊处理见 aclnn_scatter_pa_kv_cache.cpp空 Tensor 短路当 key及 value若非空对应 shape 的 shapeSize 为 0 时直接返回 workspaceSize 0 的空执行器不触发实际计算。非连续 Tensor 处理当 key 为 3 维且 scatterMode 为空/None 时通过IsSupportNonContiguousCache/IsSupportNonContiguousKeyAndCache检测 cache 首轴非连续或 key/cache 非连续但末轴连续的情况分别走ProcessNonContiguous用CreateView保留视图、避免拷贝或ProcessContiguous统一l0op::Contiguous归一后再经l0op::ViewCopy回写原始 cache 视图。这也是参数表中各 Tensor 均标注支持非连续的实现基础。第二段接口则通过CommonOpExecutorRun将执行器在指定 stream 上发射见 aclnn_scatter_pa_kv_cache.cpp。参数详解下表汇总第一段接口的完整参数说明数据格式均为 ND非连续 Tensor 均支持参数名输入/输出描述使用说明数据类型维度keyaclTensor*输入待更新的 key 值当前 step 多个 token 的 key不支持空 TensorFLOAT16、FLOAT、BFLOAT16、INT8、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT4_E1M2、FLOAT4_E2M13-4keyCacheRefaclTensor*输入/输出需要更新的 key cache当前 layer 的 key cache不支持空 TensorcacheMode 为空指针或 Norm 时仅支持 ND 排布PA_NZ 时仅支持 FRACTAL_NZ 排布与 key 保持一致4-5slotMappingaclTensor*输入每个 token key 或 value 在 cache 中的存储偏移不支持空 TensorINT32、INT641valueaclTensor*输入待更新的 value 值当前 step 多个 token 的 value不支持空 Tensor非 0 维下 shape 与 key 一致与 key 保持一致0、3、4valueCacheRefaclTensor*输入/输出需要更新的 value cache当前 layer 的 value cache不支持空 Tensor非 0 维下 shape 与 keyCacheRef 一致排布约束同 keyCacheRef与 key 保持一致0、4、5compressLensOptionalaclTensor*可选输入压缩量-与 slotMapping 保持一致1compressSeqOffsetOptionalaclTensor*可选输入每个 batch 每个 head 的压缩起点-与 slotMapping 保持一致1seqLensOptionalaclTensor*可选输入每个 batch 的实际 seqLens-与 slotMapping 保持一致1cacheModeOptionalchar*输入表示 keyCacheRef 和 valueCacheRef 的内存排布格式空指针或 Norm → NDPA_NZ → FRACTAL_NZ--scatterModeOptionalchar*输入表示更新的 key 和 value 的状态空指针或 None非压缩且连续AlibiAlibi 结构压缩RopeRope 结构压缩OmniOmni 结构压缩Nct非压缩但非连续NHSDcache 格式为 [num_blocks, num_head, block_size, head_size]--stridesOptionalaclIntArray*输入key 和 value 在非连续状态下的步长数组长度 2值应大于 0仅 scatterMode 为 Nct 时生效分别表示 strideK、strideV--offsetsOptionalaclIntArray*输入key 和 value 在非连续状态下的偏移数组长度 2值应大于等于 0仅 scatterMode 为 Nct 时生效分别表示 offsetK、offsetV--workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小---executoraclOpExecutor**输出返回 op 执行器包含算子计算流程---产品相关的数据类型限制Atlas A2 / A3 系列输入 key、keyCacheRef、value、valueCacheRef 仅支持 FLOAT16、BFLOAT16、INT8不支持 FLOAT、UINT8、INT16、UINT16、INT32、UINT32、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN、FLOAT4_E1M2、FLOAT4_E2M1。这与 aclnn_scatter_pa_kv_cache.cpp 中的KEY_VALUE_DTYPE_SUPPORT_LIST_910仅 FLOAT16、BF16、INT8以及 scatter_pa_kv_cache_def.cpp 中keyDataType的定义一致。Ascend 950PR / 950DT支持参数表中列出的全部数据类型其中 FLOAT4_E1M2、FLOAT4_E2M1 仅在场景一、场景二且 scatterMode 为 None 时支持。对应 scatter_pa_kv_cache_def.cpp 中的DataType_950ascend950/ascend350 配置。Rope / Omni 模板tiling 侧额外限制仅支持 FLOAT16、BFLOAT16见 scatter_pa_kv_cache_tiling.cpp。返回值与错误码两段接口均返回aclnnStatus状态码具体参见 aclnn 返回码说明。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 key、keyCacheRef、slotMapping、value、valueCacheRef 是空指针ACLNN_ERR_PARAM_INVALID161002参数 key、value、slotMapping 的数据类型不在支持范围内ACLNN_ERR_PARAM_INVALID161002key 的维数不等于 3 维或 4 维value 的维数不等于 0 维、3 维或 4 维ACLNN_ERR_PARAM_INVALID161002key、keyCacheRef、value、valueCacheRef 的数据类型不一致ACLNN_ERR_PARAM_INVALID161002slotMapping、compressLensOptional、compressSeqOffsetOptional、seqLensOptional 的数据类型不一致上述校验可在 aclnn_scatter_pa_kv_cache.cpp空指针、CheckShape维度与 CheckDtypeValid类型一致性中找到对应实现。此外图模式下 inferShape 阶段还会校验cache_mode与scatter_mode的取值合法性仅允许文档列出的枚举值见 scatter_pa_kv_cache_infershape.cpp。第二段接口aclnnScatterPaKvCache的入参为 workspaceDevice 侧申请的内存地址、workspaceSize由第一段接口获取、executorop 执行器与 stream指定执行任务的 Stream。约束说明确定性计算aclnnScatterPaKvCache 默认确定性实现。数据类型key、value、keyCacheRef、valueCacheRef 的数据类型必须一致当 cacheModeOptional 为 PA_NZ 时key、keyCacheRef 与 value、valueCacheRef 的数据类型可以不一致slotMapping、compressLensOptional、compressSeqOffsetOptional、seqLensOptional 的数据类型必须一致。slotMapping 值域取值范围[0, num_blocks*block_size-1]且元素值保证不重复重复时不保证正确性。shape 一致性当 key 和 value 都是 3 维则 key 和 value 的前两维 shape 必须相同当 key 和 value 都是 4 维则前三维 shape 必须相同且 keyCacheRef 和 valueCacheRef 的第三维必须是 1当 key 和 value 是 4 维时compressLensOptional、seqLensOptional 为必选参数3 维时三者均为可选参数4 维场景场景三下 slotMapping 为 2 维第一维等于 key 的第一维 batch第二维等于 key 的第三维 num_headseqLensOptional 为 1 维且长度等于 batch3 维且存在 seqLensOptional 时场景四、五seqLensOptional 中所有值的和等于 key 的第一维 num_tokens场景三、四、五中seqLensOptional 和 compressLensOptional 的每个元素必须满足reduceSum(seqLensOptional[i] - compressLensOptional[i]) num_blocks * block_size。场景一PA_NZkeyCacheRef 和 valueCacheRef 的倒数第二维必须小于 UINT16_MAX。场景七NHSDk_head_size 和 v_head_size 必须 32 字节对齐num_head 必须小于 4095对应 tiling 中的NUM_HEAD_MAX 4095datacopypad 基础接口跳搬块数上限block_size * k_head_size和block_size * v_head_size必须小于 UINT32_MAX(num_head * k_head_size num_head * v_head_size) * sizeof(dtype_key)必须小于 196608即单个 token 的 key 加 value 大小不能超过 UB 容量校验见 scatter_pa_kv_cache_tiling.cpp。调用示例以下示例展示完整的 aclnn 两段式调用流程对应场景一PA_NZ 模式完整可编译工程可参考 examples/test_aclnn_scatter_pa_kv_cache.cpp具体编译与执行过程请参见 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_scatter_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, aclFormat format) { 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, format, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.固定写法device/stream初始化参考acl API手册 // 根据自己的实际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的接口自定义构造 // 场景一示例key[2, 2, 32]keyCache[1, 4, 32, 16]PA_NZ 模式 std::vectorint64_t keyShape {2, 2, 32}; std::vectorint64_t keyCacheShape {1, 4, 32, 16}; std::vectorint64_t slotMappingShape {2}; std::vectorint64_t valueShape {2, 2, 32}; std::vectorint64_t valueCacheShape {1, 4, 32, 16}; void* keyDeviceAddr nullptr; void* valueDeviceAddr nullptr; void* slotMappingDeviceAddr nullptr; void* keyCacheDeviceAddr nullptr; void* valueCacheDeviceAddr nullptr; void* compressLensDeviceAddr nullptr; void* compressSeqOffsetDeviceAddr nullptr; void* seqLensDeviceAddr nullptr; aclTensor* key nullptr; aclTensor* value nullptr; aclTensor* slotMapping nullptr; aclTensor* keyCache nullptr; aclTensor* valueCache nullptr; aclTensor* compressLens nullptr; aclTensor* compressSeqOffset nullptr; aclTensor* seqLens nullptr; char * cacheMode PA_NZ; char * scatterMode None; std::vectorint16_t hostKey(128, 0); std::vectorint16_t hostValue(128, 0); std::vectorint32_t hostSlotMapping {0, 1}; std::vectorint16_t hostKeyCacheRef(2048, 1); std::vectorint16_t hostValueCacheRef(2048, 1); std::vectorint64_t hostStrides(2, 1); std::vectorint64_t hostOffsets(2, 0); // 创建key aclTensor ret CreateAclTensor(hostKey, keyShape, keyDeviceAddr, aclDataType::ACL_FLOAT16, key, aclFormat::ACL_FORMAT_ND); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建value aclTensor ret CreateAclTensor(hostValue, valueShape, valueDeviceAddr, aclDataType::ACL_FLOAT16, value, aclFormat::ACL_FORMAT_ND); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建slotMapping aclTensor ret CreateAclTensor(hostSlotMapping, slotMappingShape, slotMappingDeviceAddr, aclDataType::ACL_INT32, slotMapping, aclFormat::ACL_FORMAT_ND); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建keyCache aclTensor ret CreateAclTensor(hostKeyCacheRef, keyCacheShape, keyCacheDeviceAddr, aclDataType::ACL_FLOAT16, keyCache, aclFormat::ACL_FORMAT_ND); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建valueCache aclTensor ret CreateAclTensor(hostValueCacheRef, valueCacheShape, valueCacheDeviceAddr, aclDataType::ACL_FLOAT16, valueCache, aclFormat::ACL_FORMAT_ND); CHECK_RET(ret ACL_SUCCESS, return ret); aclIntArray *strides aclCreateIntArray(hostStrides.data(), 2); CHECK_RET(strides ! nullptr, return ACL_ERROR_INTERNAL_ERROR); aclIntArray *offsets aclCreateIntArray(hostOffsets.data(), 2); CHECK_RET(offsets ! nullptr, return ACL_ERROR_INTERNAL_ERROR); // 3. 调用CANN算子库API需要修改为具体的Api名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnScatterPaKvCache第一段接口 ret aclnnScatterPaKvCacheGetWorkspaceSize(key, keyCache, slotMapping, value, valueCache, compressLens, compressSeqOffset, seqLens, cacheMode, scatterMode, strides, offsets, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnScatterPaKvCacheGetWorkspaceSize 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); } // 调用aclnnScatterPaKvCache第二段接口 ret aclnnScatterPaKvCache(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnScatterPaKvCache 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 keyCacheSize GetShapeSize(keyCacheShape); std::vectorint16_t keyCacheResult(keyCacheSize, 0); ret aclrtMemcpy(keyCacheResult.data(), keyCacheResult.size() * sizeof(int16_t), keyCacheDeviceAddr, keyCacheSize * sizeof(int16_t), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy keyCache from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i keyCacheSize; i) { LOG_PRINT(keyCache[%ld] is: %d\n, i, keyCacheResult[i]); } auto valueCacheSize GetShapeSize(valueCacheShape); std::vectorint16_t valueCacheResult(valueCacheSize, 0); ret aclrtMemcpy(valueCacheResult.data(), valueCacheResult.size() * sizeof(int16_t), valueCacheDeviceAddr, valueCacheSize * sizeof(int16_t), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy valueCache from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i valueCacheSize; i) { LOG_PRINT(valueCache[%ld] is: %d\n, i, valueCacheResult[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(key); aclDestroyTensor(value); aclDestroyTensor(slotMapping); aclDestroyTensor(keyCache); aclDestroyTensor(valueCache); aclDestroyTensor(compressLens); aclDestroyTensor(compressSeqOffset); aclDestroyTensor(seqLens); aclDestroyIntArray(strides); aclDestroyIntArray(offsets); // 7. 释放device资源需要根据具体API的接口定义参数 aclrtFree(keyDeviceAddr); aclrtFree(valueDeviceAddr); aclrtFree(slotMappingDeviceAddr); aclrtFree(keyCacheDeviceAddr); aclrtFree(valueCacheDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点解读示例构造的是场景一PA_NZcacheMode PA_NZkeyCache/valueCache 形状为[1, 4, 32, 16]FLOAT16 下 last_dim 32/2 16num_head * k_head_size // last_dim_k 2*32/16 4满足 4 维紧凑排布的 shape 约束。strides/offsets虽在场景一中不生效但接口签名要求传入示例中分别默认填充{1, 1}与{0, 0}与算子属性默认值一致见 scatter_pa_kv_cache_def.cpp。三个可选参数compressLens、compressSeqOffset、seqLens在场景一中传nullptr即可。从接口到底层算子执行链路了解调用链有助于排查问题与评估性能特征aclnn 封装层aclnn_scatter_pa_kv_cache.cpp 完成空指针、类型、维度校验并将非连续输入归一为连续视图后调用底层l0op::ScatterPaKvCache构建执行器。图模式 / GE 侧图模式调用可参考 test_geir_scatter_pa_kv_cache.cpp 与算子 IR 定义 scatter_pa_kv_cache_proto.hscatter_pa_kv_cache_infershape.cpp 负责 shape/dataType 推导输出 keyCache/valueCache 与输入 shape 一致输出数据类型分别取自 key 与 value。tiling 层scatter_pa_kv_cache_tiling.cpp 根据 cacheMode/scatterMode 映射模板类型NORMAL1、NZ2、ALIBI3、ROPE4、SISO5、OMNI6、NORM_NCT7、SISO_NCT8、NHSD9依据 AIV 核数与 token 数计算blockFactor/usedCoreNum/tailBlockFactor并给出两种 tiling keytemplateType * 1000 1fully load整个 head 维度可一次装入 UB或templateType * 1000not fully load需分块搬运。Rope/Omni 模板的 KV 分核任务量翻倍TASK_MULTIPLE 2。kernel 层scatter_pa_kv_cache.cpp 按 tiling key 分发到 Normal/NZ/Alibi/Rope/Omni/SISO/NHSD 等具体实现其中 Rope、Omni、NHSD 分支仅在 FLOAT16/BFLOAT16 下编译生效。arch35 平台的专用实现位于 op_kernel/arch35按 fully load / not fully load / 非连续等维度拆分了独立的头文件实现。测试覆盖仓库在 tests/ut/op_host 下提供 inferShape 与 tiling 的单元测试test_scatter_pa_kv_cache_infershape.cpp、test_scatter_pa_kv_cache_tiling.cpp并在 tests/ut/op_kernel 下提供 kernel 级测试与数据生成脚本gen_data.py、gen_tiling.py可作为理解各场景 shape 与数据布局的参考。延伸阅读算子在仓库中的模块总览与图模式调用说明attention/scatter_pa_kv_cache/README.md两段式接口规范docs/zh/context/two_phase_api.mdaclnn 返回码说明docs/zh/context/aclnn_return_code.md编译与运行样例docs/zh/context/compile_and_run_sample.md非连续 Tensor 支持说明docs/zh/context/non_contiguous_tensor.md【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考