aclnnDequantRopeQuantKvcache 算子深度解析：NPU 上融合反量化、RoPE 旋转位置编码与量化 KV Cache 写入的实战指南

aclnnDequantRopeQuantKvcache 算子深度解析NPU 上融合反量化、RoPE 旋转位置编码与量化 KV Cache 写入的实战指南【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读aclnnDequantRopeQuantKvcache是 CANN ops-transformer 算子库中面向大模型推理场景的高阶融合算子它将反量化Dequant→ QKV 切分 → 旋转位置编码RoPE→ 量化Quant→ KV Cache 写入更新五个步骤融合为一次 NPU 算子调用避免了多个中间张量的 Device↔Host 往返与多次 Kernel 启动开销。阅读本文后你将掌握该算子的完整计算流程、两段式 aclnn 接口的每个参数语义与取值约束、两种 KV Cache 更新模式contiguous / page的差异并能够依据仓库中的完整调用示例与源码实现在自己的推理管线中正确接入该算子。本算子位于仓库的 posembedding/dequant_rope_quant_kvcache 目录完整源码可参考该目录下的 op_host 与 op_kernel 实现。产品支持情况根据算子官方文档与目录内 README.md 说明当前支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√需要特别说明两点Kirin 平台不支持 BFLOAT16README 明确标注因此在 Kirin X90 / Kirin 9030 上使用本算子时x 仅支持 FLOAT16 / INT32 两种类型。这一差异在算子注册源码中也有体现dequant_rope_quant_kvcache_def.cpp 为 Kirin 平台单独定义了XDtypeListKirin/cosDtypeListKirin等数据类型列表其中不含DT_BF16且 AICore 配置kirinx90、kirin9030通过GetKirinCoreConfig()统一挂载。在 dequant_rope_quant_kvcache_def.cpp 中可以看到算子通过this-AICore().AddConfig(ascend910b)、AddConfig(ascend910_93)、AddConfig(ascend950)完成 AICore 配置注册与上表支持的产品一一对应。算子功能与计算流程功能总览算子对输入张量x执行如下流水线Dequant可选对输入x进行反量化恢复高精度浮点表示切分Split按属性sizeSplits给出的长度对尾轴dim-1进行切分得到q、k、vOut三段RoPE旋转位置编码对q、k应用基于cos、sin的旋转位置编码生成qOut和kOutQuant量化对kOut与vOut分别使用scaleK/offsetK、scaleV/offsetV进行静态量化输出 INT8 数据KV Cache 更新根据indices指定的 token 位置信息将量化后的 k、v 写入kCacheRef与vCacheRef。从计算路径看该算子本质上是把 LLM 推理中权重反量化 QKV 投影结果切分 位置编码 KV 量化缓存这段高频热点串成一条端到端的单算子流水适合与 PagedAttention 类推理框架配合使用。计算步骤公式算子的计算过程可形式化为以下五步第 1 步反量化可选$$ dequantX Dequant(x, weightScaleOptional, activationScaleOptional, biasOptional) $$第 2 步尾轴切分$$ q, k, vOut SplitTensor(dequantX, dim-1, sizeSplits) $$第 3 步旋转位置编码$$ qOut, kOut ApplyRotaryPosEmb(q, k, cos, sin) $$第 4 步静态量化$$ quantK Quant(kOut, scaleK, offsetKOptional) $$$$ quantV Quant(vOut, scaleV, offsetVOptional) $$第 5 步KV Cache 更新两种模式见下文两种 KV Cache 更新模式模式一cacheModeOptional contiguous默认连续式缓存按 batch 维度逐位置写入$$ kCacheRef[i][indices[i]] quantK[i] $$$$ vCacheRef[i][indices[i]] quantV[i] $$模式二cacheModeOptional page分页式缓存Paged KV Cache先将 4 维 cache 张量展平成[总页数, 页内行数, 列数]的视图再按下标写入$$ kCacheRefView kCacheRef.view(-1, kCacheRef[-2], kCacheRef[-1]) $$$$ vCacheRefView vCacheRef.view(-1, vCacheRef[-2], vCacheRef[-1]) $$$$ kCacheRefView[indices[i]] quantK[i] $$$$ vCacheRefView[indices[i]] quantV[i] $$两种模式的选择会直接影响indices的 shape 语义见参数表与约束说明并且 tiling 阶段会通过CheckPaCacheMode()见 dequant_rope_quant_kvcache_tiling.cpp识别是否为 page 模式进而把batch展开为B*S、seqlen置 1 参与任务划分。两段式接口与函数原型该算子采用 CANN 算子库标准的两段式接口详见仓库文档 两段式接口说明必须先调用aclnnDequantRopeQuantKvcacheGetWorkspaceSize接口获取入参校验结果与所需 workspace 大小再调用aclnnDequantRopeQuantKvcache接口执行计算。第一段接口原型aclnnStatus aclnnDequantRopeQuantKvcacheGetWorkspaceSize( const aclTensor *x, const aclTensor *cos, const aclTensor *sin, aclTensor *kCacheRef, aclTensor *vCacheRef, const aclTensor *indices, const aclTensor *scaleK, const aclTensor *scaleV, const aclTensor *offsetKOptional, const aclTensor *offsetVOptional, const aclTensor *weightScaleOptional, const aclTensor *activationScaleOptional, const aclTensor *biasOptional, const aclIntArray *sizeSplits, char *quantModeOptional, char *layoutOptional, bool kvOutput, char *cacheModeOptional, const aclTensor *qOut, const aclTensor *kOut, const aclTensor *vOut, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnDequantRopeQuantKvcache( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)第一段接口完成入参校验与 workspace 大小计算第二段接口在指定stream上真正下发执行。两段接口均返回aclnnStatus状态码具体取值参见仓库文档 aclnn 返回码说明。参数详解第一段接口参数表下表完整列出aclnnDequantRopeQuantKvcacheGetWorkspaceSize的全部参数语义。其中非连续 Tensor标记为 √ 表示该输入支持非连续带 stride张量。参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx输入公式中用于切分的输入 xshape 为[B, S, H]或[B, H]H(NqNkvNkv)*D。x 的尾轴小于等于 4096且按 64 对齐FLOAT16、BFLOAT16、INT32ND2-3√cos输入公式中用于位置编码的输入 cosx 为 3 维时 shape 为[B, S, 1, D]x 为 2 维时 shape 为[B, D]FLOAT16、BFLOAT16ND24√sin输入公式中用于位置编码的输入 sinx 为 3 维时 shape 为[B, S, 1, D]x 为 2 维时 shape 为[B, D]和 cos 保持一致ND24√kCacheRef输入公式中用于缓存 k 的输入 kCacheRefshape 为[C_1, C_2, Nkv, D]INT8ND4√vCacheRef输入公式中用于缓存 v 的输入 vCacheRefshape 为[C_1, C_2, Nkv, D]INT8ND4√indices输入表示 Kvcache 的 token 位置信息的输入 indices当 cache_mode 为 page 且 x 为 3 维时 shape 为[B*S]否则 shape 为[B]INT32ND1√scaleK输入公式中的输入 scaleK用于量化 k 的 scale 因子元素个数为Nkv*D推荐 shape 为[Nkv, D]兼容一维展平 shape[Nkv*D]FLOATND≥1√scaleV输入公式中的输入 scaleV用于量化 v 的 scale 因子元素个数为Nkv*D推荐 shape 为[Nkv, D]兼容一维展平 shape[Nkv*D]FLOATND≥1√offsetKOptional输入公式中的输入 offsetKOptional用于量化 k 的 offset 因子元素个数为Nkv*D推荐 shape 为[Nkv, D]兼容一维展平 shape[Nkv*D]FLOATND≥1√offsetVOptional输入公式中的输入 offsetVOptional用于量化 v 的 offset 因子元素个数为Nkv*D推荐 shape 为[Nkv, D]兼容一维展平 shape[Nkv*D]FLOATND≥1√weightScaleOptional输入公式中的输入 weightScaleOptional用于反量化的权重 scale 因子shape 为[H]FLOATND1√activationScaleOptional输入公式中的输入 activationScaleOptional用于反量化的激活 scale 因子x 为 3 维时 shape 为[B*S]x 为 2 维时 shape 为[B]FLOATND1√biasOptional输入公式中的输入用于反量化的偏置 biasOptionalshape 为[H]FLOAT、FLOAT16、INT32、BFLOAT16ND1√sizeSplits输入表示输入的 qkv 进行切分的长度size 大小为 3值为[Nq*D, Nkv*D, Nkv*D]AclIntArray---quantModeOptional输入表示支持的量化类型目前仅传入staticCHAR---layoutOptional输入表示支持的数据格式目前仅支持BSNDCHAR---kvOutput输入Host 侧布尔值表示是否输出 kOut 和 vOut为 true 时输出有效 shape 的 kOut 和 vOut为 false 时 kOut 和 vOut 的 shape 为空BOOL---cacheModeOptional输入表示 kCacheRef 的更新方式目前仅支持page和contiguous默认为contiguousCHAR---qOut输出公式中经旋转位置编码后的 qx 为 3 维时 shape 为[B, S, Nq, D]x 为 2 维时 shape 为[B, Nq, D]。数据类型与 cos、sin 保持一致FLOAT16、BFLOAT16ND3-4×kOut输出公式中经旋转位置编码后的 kkvOutput 为 true 时x 为 3 维时 shape 为[B, S, Nkv, D]x 为 2 维时 shape 为[B, Nkv, D]kvOutput 为 false 时 shape 为空。数据类型与 cos、sin 保持一致FLOAT16、BFLOAT16ND1-4×vOut输出公式中切分得到的 vkvOutput 为 true 时x 为 3 维时 shape 为[B, S, Nkv, D]x 为 2 维时 shape 为[B, Nkv, D]kvOutput 为 false 时 shape 为空。数据类型与 cos、sin 保持一致FLOAT16、BFLOAT16ND1-4×workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----参数语义的源码佐证上述参数定义可以在算子注册源码 dequant_rope_quant_kvcache_def.cpp 中得到印证x、cos、sin、k_cache、v_cache、indices、scale_k、scale_v为REQUIRED必选输入offset_k、offset_v、weight_scale、activation_scale、bias为OPTIONAL可选输入所有输入均声明了AutoContiguous()与参数表中非连续 Tensor的支持情况一致属性侧size_splits为 REQUIRED 的ListIntquant_mode默认static、layout默认BSND、kv_output默认false、cache_mode默认contiguous与文档描述完全对应输入/输出均限定FORMAT_ND格式即文档参数表中的数据格式 ND。第二段接口参数表参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnDequantRopeQuantKvcacheGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值与错误码两段接口均返回aclnnStatus状态码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入和输出的 Tensor 是空指针ACLNN_ERR_PARAM_INVALID161002输入和输出的数据类型不在支持的范围内完整返回码语义参见仓库文档 aclnn 返回码说明。约束说明使用该算子时必须满足以下约束确定性计算aclnnDequantRopeQuantKvcache默认确定性实现。contiguous 模式下的 indices 取值kCacheRef的第 0 维大于等于 x 的第 0 维。x 为 3 维时indices 数据值大于等于 0 且小于等于kCacheRef的第 1 维减 x 的第 1 维x 为 2 维时indices 数据值大于等于 0 且小于等于kCacheRef的第 1 维减 1。page 模式下的 indices 取值indices 数据值大于等于 0小于kCacheRef的第 0 维 × 第 1 维且不重复。非 INT32 输入输入 x 不为 INT32 时x、cos、sin 与输出 qOut、kOut、vOut 的数据类型保持一致此时activationScaleOptional、weightScaleOptional、biasOptional不生效。INT32 输入反量化路径输入 x 为 INT32 时cos、sin 与输出 qOut、kOut、vOut 的数据类型保持一致此时weightScaleOptional必选activationScaleOptional、biasOptional可选biasOptional不需要与其他输入类型一致。尾轴限制x 的尾轴小于等于 4096且按 64 对齐。Kirin 平台Kirin X90 / Kirin 9030 处理器系列产品不支持 BFLOAT16。这些约束在 tiling 阶段会被进一步强制校验。例如 dequant_rope_quant_kvcache_tiling.cpp 中会检查sizeSplits长度必须为 3、k 与 v 的切分长度必须相等、hiddenSize必须为 16 的倍数、qHiddenSize/vHiddenSize必须为hiddenSize的整数倍由此推出 Nq、Nkv 必须为整数、x 的尾轴必须等于三段切分长度之和等同时 tiling 还会校验 scale/offset 的元素个数必须等于Nkv*DquantShapeSize与参数表要求一致。调用示例仓库在 examples/test_aclnn_dequant_rope_quant_kvcache.cpp 中提供了可直接参考的完整示例其调用流程与本算子文档中的示例代码一致编译和执行过程请参考仓库文档 编译与运行样例。核心代码如下#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_dequant_rope_quant_kvcache.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; } void PrintOutResult(std::vectorint64_t shape, void** deviceAddr) { auto size GetShapeSize(shape); std::vectorint8_t resultData(size, 0); auto ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, 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); for (int64_t i 0; i size; i) { LOG_PRINT(mean result[%ld] is: %d\n, i, resultData[i]); } } 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 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的接口定义构造 int64_t shapeB 1; // batch int64_t shapeS 1; // seqlen int64_t shapeNq 2; // query 头数 int64_t shapeNkv 1; // kv 头数 int64_t shapeD 32; // 每头维度 int64_t shapeH shapeD * (shapeNq shapeNkv shapeNkv); // 尾轴长度 (NqNkvNkv)*D std::vectorint64_t inputShape {shapeB, shapeS, shapeH}; std::vectorint64_t cosShape {shapeB, shapeS, 1, shapeD}; std::vectorint64_t sinShape {shapeB, shapeS, 1, shapeD}; std::vectorint64_t kcacheShape {shapeB, shapeH, 1, shapeD}; std::vectorint64_t vcacheShape {shapeB, shapeH, 1, shapeD}; std::vectorint64_t indicesShape {shapeB}; std::vectorint64_t kscaleShape {shapeNkv, shapeD}; std::vectorint64_t vscaleShape {shapeNkv, shapeD}; std::vectorint64_t koffsetShape {shapeNkv, shapeD}; std::vectorint64_t voffsetShape {shapeNkv, shapeD}; std::vectorint64_t weightShape {shapeH}; std::vectorint64_t activationShape {shapeB * shapeS}; std::vectorint64_t biasShape {shapeH}; // 以 INT32 输入触发反量化路径需提供 weightScale/activationScale/bias std::vectorint32_t inputHostData(shapeB * shapeS * shapeH, 1); std::vectorint16_t cosHostData(shapeB * shapeS * shapeD, 1); std::vectorint16_t sinHostData(shapeB * shapeS * shapeD, 1); std::vectorint8_t kcacheHostData(shapeB * shapeH * shapeD, 6); std::vectorint8_t vcacheHostData(shapeB * shapeH * shapeD, 6); std::vectorint32_t indicesHostData(shapeB, 0); std::vectorfloat kscaleHostData(shapeNkv * shapeD, 2); std::vectorfloat vscaleHostData(shapeNkv * shapeD, 2); std::vectorfloat koffsetHostData(shapeNkv * shapeD, 2); std::vectorfloat voffsetHostData(shapeNkv * shapeD, 2); std::vectorfloat weightHostData(shapeH, 2); std::vectorfloat activationHostData(shapeB * shapeS, 2); std::vectorfloat biasHostData(shapeH, 2); // 省略逐个调用 CreateAclTensor 创建各输入/输出的 aclTensor 与 device 内存 std::vectorint64_t splitData {shapeNq * shapeD, shapeNkv * shapeD, shapeNkv * shapeD}; aclIntArray *sizeSplits aclCreateIntArray(splitData.data(), splitData.size()); char quantMode[] static; char layout[] BSND; char cacheMode[] contiguous; // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnDequantRopeQuantKvcache第一段接口 ret aclnnDequantRopeQuantKvcacheGetWorkspaceSize(input, cos, sin, kcache, vcache, indices, kscale, vscale, koffset, voffset, weight, activation, bias, sizeSplits, quantMode, layout, true, cacheMode, q, k, v, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDequantRopeQuantKvcacheGetWorkspaceSize 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); } // 调用aclnnDequantRopeQuantKvcache第二段接口 ret aclnnDequantRopeQuantKvcache(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDequantRopeQuantKvcache 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侧 PrintOutResult(kcacheShape, kcacheDeviceAddr); PrintOutResult(vcacheShape, vcacheDeviceAddr); // 6. 释放aclTensor和aclIntArray // 7. 释放device资源aclrtFree / aclrtDestroyStream / aclrtResetDevice / aclFinalize return 0; }上述示例中sizeSplits {Nq*D, Nkv*D, Nkv*D}与参数表要求一致。需要注意示例中indices全部为 0 且cacheMode contiguous因此量化后的 k/v 会被写入每一条 batch 对应 cache 的第 0 个 token 位置若要在真实推理中做增量写入需按 token 位置更新indices。源码级原理深挖算子定义与数据类型注册dequant_rope_quant_kvcache_def.cpp 中通过OP_ADD(DequantRopeQuantKvcache)完成算子注册。该文件用 16 组 dtype 组合FLOAT16 / INT32 / BF16 三种 x 类型分别与 cos/sin、bias、scale、cache 类型的组合精确刻画了文档约束中数据类型保持一致的规则例如x为 FLOAT16 时cos/sin必须为 FLOAT16XDtypeList与cosDtypeList前 4 组一一对应x为 INT32 时bias可以是 FLOAT16 / BF16 / INT32 / FLOAT 四种biasDtypeList中对应条目即约束第 5 条biasOptional 不需要与其他输入类型一致k_cache/v_cache固定为 INT8indices固定为 INT32scale/offset/weight_scale/activation_scale固定为 FLOAT。Shape 推导逻辑dequant_rope_quant_kvcache_infershape.cpp 中InferShapeForDequantRopeQuantKvcache负责输出 shape 推导要求 x 为 2 维或 3 维cache 为 4 维从size_splits[0]与 cache 的隐藏维度推导qHead size_splits[0] / hiddenSizex 为 3 维时qOut为[B, S, Nq, D]kOut/vOut为[B, S, Nkv, D]x 为 2 维时seqlen 视为 1qOut为[B, Nq, D]kOut/vOut为[B, Nkv, D]当kv_output为 false 时kOut/vOut的第 0 维被置为 0shape 为空这与参数表中kvOutput 为 false 时 shape 为空完全对应InferDataTypeForDequantRopeQuantKvcache将 q/k/v 的输出类型设置为与cos一致印证数据类型与 cos、sin 保持一致。Tiling 策略与 workspacedequant_rope_quant_kvcache_tiling.cpp 中的TilingDequantRopeQuantKvcache负责在 Host 侧完成任务切分任务量建模taskNum batch * seqlen并按 AIV 核数GetCoreNumAiv()与 UB 容量GetCoreMemSize计算blockFactor前核块因子与tailCoreBlockFactor尾核块因子实现多核负载均衡UB 单次承载量OnceUBMaxS由 UB 剩余空间除以单次处理所需的 buffer 总量q/k/v/cos/sin/indices 的缓冲区之和全部按 32 字节BLOCK_SIZE对齐计算得到Kernel 按该值循环搬数Workspacetiling 阶段统一申请MINIMAL_WORKSPACE 16MB的 workspace见 dequant_rope_quant_kvcache_tiling.cpp作为空 tensor 与 kernel 计算时的中间缓冲TilingKey以 bias 的数据类型作为SetTilingKey的取值FLOAT0 / FLOAT161 / INT322 / BF163驱动 Kernel 侧模板实例化。tiling 数据的字段定义位于 dequant_rope_quant_kvcache_tiling.hqHeadNum、kvHeadNum、hiddenSize、OnceUBMaxS、isPA、ifKVout、hasBias、hasAS等这些字段直接决定了 Kernel 内部的分支走向。Kernel 实现要点Kernel 侧入口为 dequant_rope_quant_kvcache.cpp按 TilingKey 0/1/2/3 分别实例化RopeQuantKvcacheV2DTYPE_X, bias类型, DTYPE_COS模板bias 类型对应 FLOAT / half / int32_t / bfloat16_t。核心类实现在 dequant_rope_quant_kvcache.h反量化dequantUb对 INT32 输入先Cast到 float再乘weight_scaleMul可选乘激活 scaleMuls与加 biasAdd切分搬数通过DataCopyPad与DataCopyExtParams以块内 stride方式从inputGm中带间隔地抽取 q、k、v 三段见dataCopyParamsQ_/K_/V_的构造srcStride恰好跳过其他两段避免三块独立访存RoPE 计算按hiddenSize / 2拆分奇偶半段执行k*cos ± 旋转(k)*sin的旋转位置编码代码中对sin先乘以 -1Muls再通过两次Mul与一次Add完成标准 RoPE 公式量化对 kOut/vOut 先Div除以 scale可选Addoffset再经Cast(CAST_RINT)到 INT16、转 half、最终Cast(CAST_NONE)到 INT8Cache 写入copyOutcachepage 模式下以index * kvHeadNum * hiddenSize计算页内偏移contiguous 模式下以(bOffset bIndex) * cacheSeqlen index sIndex计算连续偏移随后DataCopy写回kCacheGm/vCacheGm流水并行通过MTE2_S、V_MTE3、MTE3_MTE2、MTE2_MTE3等硬件事件SetFlag/WaitFlag对搬入MTE2、向量计算V、搬出MTE3三段流水做同步编排降低访存延迟。测试与验证仓库为算子提供了多级测试保障ST 测试目录 tests/st/aclnnDequantRopeQuantKvcache 下的atk_aclnnDequantRopeQuantKvcache.json定义了基于 ATK 的端到端用例用例 0 覆盖 x 为[1, 2304]的 FLOAT16 输入Nq*D1536、Nkv*D384、Nkv*D384即 2D 输入 cacheModepage用例 1 覆盖 x 为[1, 3584]的 INT32 反量化路径sizeSplits{1792, 896, 896}bias 为 BF16cacheModepage。两个用例均携带backward: true可用于精度对比基准cv_fused_double_benchmarkUT 测试目录 tests/ut 下包含 op_host 层的test_dequant_rope_quant_kvcache_infershape.cpp、test_dequant_rope_quant_kvcache_tiling.cpp以及 op_kernel 层的test_dequant_rope_quant_kvcache.cpp分别验证 shape 推导、tiling 数据与 Kernel 计算结果的正确性。小结aclnnDequantRopeQuantKvcache是 CANN ops-transformer 中把反量化 → QKV 切分 → RoPE → 量化 → KV Cache 写入五步融合的单算子实现。通过两段式 aclnn 接口开发者可以在 Host 侧一次性完成参数校验、workspace 计算与执行器构建随后在 Stream 上异步执行。本文完整梳理了其计算流程、两段式接口的全部参数语义、两种 cache 更新模式与 7 条使用约束并结合仓库中的算子定义、shape 推导、tiling 策略与 Kernel 实现揭示了底层的数据类型组合规则、多核任务划分与流水并行机制。对于正在自研或接入大模型推理框架、需要在 NPU 上高效维护量化 KV Cache 的开发者该算子可作为 PagedAttention 场景中 KV 预处理环节的落地参考。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考