CANN ops-transformer aclnnKvCompressEpilog 算子详解KV Cache Epilog 阶段原地量化压缩更新【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读aclnnKvCompressEpilog 是 CANN ops-transformer 开源算子库中用于 KV Cache 压缩更新的核心算子在大模型推理的 Epilog 阶段将bfloat16激活向量按slotMapping散写scatter到分页 KV Cache 并原地压缩支持 FP8 分组动态量化与 rope/nope 分段混合量化三种模式。本文基于 aclnnKvCompressEpilog.md 与仓库源码完整讲解其功能语义、两段式接口用法、参数约束与源码实现使读者能够独立完成算子接入、参数配置与结果校验。产品支持情况该算子为 Ascend 950 系列专属特性当前仅支持以下产品其余产品均不支持产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持该结论与 op_host/kv_compress_epilog_def.cpp 中AICore().AddConfig(ascend950)的注册配置一致即算子仅在 ascend950 架构上编译与运行。功能说明算子语义在 KV Cache 的 Epilog 阶段对 KV Cache 进行原地压缩更新将bfloat16激活值x量化压缩后按slotMapping散写到 cacheslotMapping中值为-1的 token 跳过不处理x尾轴d的后 64 列为 rope 段、前d-64列为 nope 段cache 为原地更新in-place未被命中的行保留原值。支持三种量化模式quantMode0groupbf16 scale逐组动态量化压缩为 FP8scale存储为bfloat16rope 段保留bfloat16quantMode1groupe8m0 scale量化方法同模式 0但scale存储为float8_e8m02 的幂指数即 MX-FP8 微缩放语义roundScaletrue时scale向上取到 2 的幂quantMode2rope hifloat8 静态 nope FLOAT4_E2M1 动态rope 段后 64 列做 hifloat8静态量化乘xScalenope 段前d-64列做 per-groupFLOAT4_E2M1 动态量化FP4_MAX6.0scale以bfloat16写出roundScale在该模式下不生效。计算公式对x的最后一维d 轴按组计算每组的 amax 并量化为目标 dtype记第 g 组为 $x_g$场景 1quantMode0group(64) 量化为 FP8scale存储为bfloat16rope 段保留bfloat16。$$ scale_g \frac{\max(|x_g|)}{FP8_MAX}, \quad q_i \mathrm{round}\left(\frac{x_i}{scale_g}\right) $$场景 2quantMode1同场景 1但scale存储为float8_e8m0roundScaletrue时scale向上取到 2 的幂。场景 3quantMode2rope 段后 64 列做 hifloat8 静态量化nope 段前 d-64 列做 per-group FLOAT4_E2M1 动态量化FP4_MAX6.0scale以bfloat16写出roundScale不生效。$$ rope_i \mathrm{hifloat8}(x_i \cdot xScale), \qquad scale_g \frac{\max(|x_g|)}{FP4_MAX}, \quad nope_i \mathrm{FLOAT4_E2M1}\left(\frac{x_i}{scale_g}\right) $$仓库 kernel 侧实现op_kernel/kv_compress_epilog_common.h中定义了相关量化常量可作为公式的代码级印证constexpr float FP8_E5M2_MAX_VALUE 57344.0f; constexpr float FP8_E4M3FN_MAX_VALUE 448.0f; constexpr float FP8_E5M2_MIN_VALUE -57344.0f; constexpr float FP8_E4M3FN_MIN_VALUE -448.0f;FP8 量化默认采用 E4M3FNmax448.0同时保留了 E5M2max57344.0路径的常量定义。典型示例shape 速查cache shape: [128, 16, 1, 384] # 4D [blockNum, blockSize, 1, headDim], num_slots2048, headDim≥kvCacheCol(323) x shape: [1024, 256] slot_mapping shape: [1024] quantGroupSize 64 quantMode 1 roundScale true xScale 1.0 cache out shape: [128, 16, 1, 384] (原地更新)两段式接口调用模型每个算子分为两段式接口详见 两段式接口必须先调用第一段接口aclnnKvCompressEpilogGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用第二段接口aclnnKvCompressEpilog执行计算。aclnnStatus aclnnKvCompressEpilogGetWorkspaceSize( aclTensor *cacheRef, const aclTensor *x, const aclTensor *slotMapping, int64_t quantGroupSize, int64_t quantMode, bool roundScale, float xScale, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnKvCompressEpilog( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从源码看op_host/op_api/aclnn_kv_compress_epilog.cpp第一段接口完成参数校验后通过l0op::KvCompressEpilog构建 L0 图并返回 workspace 大小与执行器第二段接口调用CommonOpExecutorRun在指定 stream 上真正执行。workspace 由用户在 Device 侧通过aclrtMalloc申请执行完毕后由用户释放。aclnnKvCompressEpilogGetWorkspaceSize 参数说明参数明细参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorcacheRefaclTensor*输入/输出当前层的 KV Cache 向量缓存原地更新对应计算公式中的量化目标不支持空 Tensor既是输入也是输出仅支持四维 [blockNum, blockSize, 1, headDim]dim2 固定为 1UINT8ND[blockNum, blockSize, 1, headDim]×xconst aclTensor*输入待量化的激活输入对应计算公式中 x不支持空 Tensor最后一维 d 需满足 d % 64 0 且 64 d ≤ 8192BFLOAT16ND[bs, d]×slotMappingconst aclTensor*输入token 到 cache slot 的索引映射值为 -1 表示跳过该 token不支持空 Tensor维度需等于 x 的维度减 1元素取值范围为 [-1, num_slots - 1]INT32、INT64ND[bs]×quantGroupSizeint64_t输入量化分组大小默认值 64。quantMode2 时仅支持 16/32/64且要求 (d-64) % quantGroupSize 0----quantModeint64_t输入量化模式枚举值支持 0、1、20group 量化scale 存储为 bfloat161group 量化scale 存储为 float8_e8m02rope 段 hifloat8 静态量化 nope 段 FLOAT4_E2M1 动态量化scale 存储为 bfloat16。默认值 1----roundScalebool输入group 模式下是否对每组 scale 向上取到 2 的幂默认值 true。quantMode2 下不生效----xScalefloat输入quantMode2 时为 rope 段 hifloat8 静态量化的缩放系数quantMode0/1 下预留未使用默认值 1.0----workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含算子计算流程-----属性默认值的源码依据各属性的默认值在 op_host/kv_compress_epilog_def.cpp 的算子定义中登记与接口文档完全一致this-Attr(quant_group_size).AttrType(OPTIONAL).Int(64); this-Attr(quant_mode).AttrType(OPTIONAL).Int(1); this-Attr(round_scale).AttrType(OPTIONAL).Bool(true); this-Attr(x_scale).AttrType(OPTIONAL).Float(1.0f);其中 quantMode 在 tiling 阶段op_host/kv_compress_epilog_tiling_arch35.cpp会做严格校验取值必须为 0/1/2 三者之一mode2 下 quantGroupSize 必须是 16/32/64。返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001cacheRef、x、slotMapping 存在空指针ACLNN_ERR_PARAM_INVALID161002cacheRef、x 或 slotMapping 为空 Tensor对应实现见 op_host/op_api/aclnn_kv_compress_epilog.cpp先做OP_CHECK_NULL空指针检查再做IsEmpty()空 Tensor 检查。aclnnKvCompressEpilog 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnKvCompressEpilogGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样为 aclnnStatus 状态码参见 aclnn 返回码。约束说明确定性计算aclnnKvCompressEpilog 默认为确定性实现。cache 仅支持四维 shape[blockNum, blockSize, 1, headDim]num_slots blockNum × blockSize倒数第二维固定为 1不支持其他维数仅在 blockNum 维支持非连续分页场景下各 block 可不紧密排布但 block 内部须连续。headDim 约束cache 末维 headDim 须 ≥ 每行写出字节数 kvCacheCol 对齐(concatCol)concatCol 计算quantMode0/1(d-64) 128 ⌈(d-64)/64⌉ × scaleBytes其中 scaleBytes 在 mode02bf16 scale、mode11e8m0 scaleconcatCol 计算quantMode264 (d-64)/2 ((d-64)/quantGroupSize) × 2对齐规则quantMode1 不补齐kvCacheCol concatColquantMode0/2 按 32B 对齐kvCacheCol ⌈concatCol/32⌉ × 32。示例d256、quantMode1 → kvCacheCol323 → headDim ≥ 323d512、quantMode1 → kvCacheCol583 → headDim ≥ 583。slotMapping 的维度应等于 x 的维度减 1即 slotMapping 为 x 除最后一维外的所有维度展平。x 的最后一维d 轴需满足 d % 64 0 且 64 d ≤ 8192按每 64 个元素一组进行逐组量化。quantMode2 时quantGroupSize 仅支持 16/32/64且 nope 段长度 (d-64) 需能被 quantGroupSize 整除x 需为 bfloat16。slotMapping 中值为 -1 的 token 会被跳过不处理其余有效元素取值范围为 [0, num_slots - 1]且元素值应保证不重复重复时不保证结果正确性。kvCacheCol 的源码级解读上述 kvCacheCol 的计算逻辑完整体现在 tiling 实现中op_host/kv_compress_epilog_tiling_arch35.cppif (quantMode_ QUANT_MODE_HIF8_FP4) { // mode2 行布局: [rope hifloat8 64B][nope FLOAT4_E2M1 (d-64)/2 B][nope bf16 scale nGroup*2 B][pad] scaleCol_ (d_ - SLICE_SIZE) / quantGroupSize_; concatCol ROPE_HIF8_BYTES (d_ - SLICE_SIZE) / 2 scaleCol_ * FP4_SCALE_BYTES; } else { // mode0/1 行布局: [rope bf16 128B][nope fp8 (d-64)B][scale] scaleCol_ CeilDiv(d_ - 64, static_castint64_t(64)); int64_t scaleBytes (quantMode_ QUANT_MODE_GROUP_QUANT_BF16) ? 2 : 1; concatCol d_ - SLICE_SIZE SLICE_SIZE * 2 scaleCol_ * scaleBytes; } if (quantMode_ QUANT_MODE_GROUP_QUANT_E8M0) { kvCacheCol_ concatCol; // mode1 不补齐 } else { kvCacheCol_ RoundUp(concatCol, BLOCK_SIZE); // mode0/2 按 32B 对齐 }可以直观理解行内布局mode0/1 下 rope 段 64 列 bf16 占 128B、nope 段 (d-64) 列 FP8 占 (d-64)B每组附带 scalemode2 下 rope 段 hifloat8 占 64B、nope 段每个 FP4 占半字节(d-64)/2 B、每组附带 2B 的 bf16 scale。tiling 校验kvCacheCol_ kvCacheRowStride_即对应 headDim 约束同文件 L301-L304。C 调用示例以下完整示例代码出自 examples/test_aclnn_kv_compress_epilog.cpp具体编译和执行过程请参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_kv_compress_epilog.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 API手册 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2.构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t cacheShape {128, 16, 1, 384}; // 4D KV Cache [blockNum,blockSize,1,headDim]num_slots2048 std::vectorint64_t xShape {1024, 256}; // 待量化激活 std::vectorint64_t slotShape {1024}; // slot 索引映射 void* cacheDeviceAddr nullptr; void* xDeviceAddr nullptr; void* slotDeviceAddr nullptr; aclTensor* cache nullptr; aclTensor* x nullptr; aclTensor* slotMapping nullptr; std::vectoruint8_t cacheHostData(GetShapeSize(cacheShape), 0); std::vectoruint16_t xHostData(GetShapeSize(xShape), 0); // bfloat16 以 uint16 承载 std::vectorint32_t slotHostData(GetShapeSize(slotShape), 0); // 创建 cache aclTensor原地输入/输出 ret CreateAclTensor(cacheHostData, cacheShape, cacheDeviceAddr, aclDataType::ACL_UINT8, cache); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 x aclTensor ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_BF16, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 slotMapping aclTensor ret CreateAclTensor(slotHostData, slotShape, slotDeviceAddr, aclDataType::ACL_INT32, slotMapping); CHECK_RET(ret ACL_SUCCESS, return ret); // 属性 int64_t quantGroupSize 64; int64_t quantMode 1; // group 量化scale 存储为 float8_e8m0 bool roundScale true; float xScale 1.0f; // 3.调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnKvCompressEpilog第一段接口 ret aclnnKvCompressEpilogGetWorkspaceSize(cache, x, slotMapping, quantGroupSize, quantMode, roundScale, xScale, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnKvCompressEpilogGetWorkspaceSize 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); } // 调用aclnnKvCompressEpilog第二段接口 ret aclnnKvCompressEpilog(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnKvCompressEpilog 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.cache 为原地更新结果已写回 cacheDeviceAddr可按需拷回 host 查看 // 6.释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensor(cache); aclDestroyTensor(x); aclDestroyTensor(slotMapping); // 7.释放device资源 aclrtFree(cacheDeviceAddr); aclrtFree(xDeviceAddr); aclrtFree(slotDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例运行流程可归纳为七个步骤设备/流初始化 → 构造输入输出 Tensor含非连续 strides 计算→ 两段式 API 调用与 workspace 申请 → 流同步 → 结果读取cache 原地更新→ 释放 Tensor → 释放设备资源。仓库还提供了更精简的工程版示例 examples/test_geir_kv_compress_epilog.cpp 以及图模式GE接入的算子 IR 定义 op_graph/kv_compress_epilog_proto.h。Python 封装接口cann_ops_transformer.kv_compress_epilog除了 aclnn 接口仓库还提供了 PyTorch 封装torch_extension/kv_compress_epilog.pyAPI 详见 torchapi_kv_compress_epilog.mdcann_ops_transformer.kv_compress_epilog(cache, x, slot_mapping, *, quant_group_size64, quant_modefp8_e8m0, round_scaleTrue, x_scale1.0) - None要点原地语义该接口无返回值结果直接写回输入cacheschema 中标注Tensor(a!)调用后直接使用入参 cachequant_mode 为字符串大小写不敏感内部经KvCompressQuantMode枚举映射为算子侧 intfp8_bf16→0、fp8_e8m0→1、hifloat8_fp4→2见 kv_compress_epilog.py支持单算子模式与 TorchAir 图模式经 graph_convert 注册的 GE converter 下沉为 KvCompressEpilog 算子eager 与图模式结果一致。单算子模式调用示例import torch import torch_npu from cann_ops_transformer.ops import kv_compress_epilog # cache为四维 [blockNum, blockSize, 1, headDim]num_slots blockNum*blockSize # headDim须 kvCacheCold256/quant_mode1 时kvCacheCol323取head_dim384 满足 block_num, block_size, head_dim 128, 16, 384 bs, d 1024, 256 cache torch.zeros(block_num, block_size, 1, head_dim, dtypetorch.uint8).npu() x torch.randn(bs, d, dtypetorch.bfloat16).npu() slot_mapping torch.randint(0, block_num * block_size, (bs,), dtypetorch.int32).npu() kv_compress_epilog( cache, x, slot_mapping, quant_group_size64, quant_modefp8_e8m0, round_scaleTrue, x_scale1.0) print(cache.shape, cache.dtype)图模式torchair调用示例import torchair class KvCompressEpilogModel(torch.nn.Module): def forward(self, cache, x, slot_mapping): kv_compress_epilog( cache, x, slot_mapping, quant_group_size64, quant_modefp8_e8m0, round_scaleTrue, x_scale1.0) return cache model KvCompressEpilogModel().npu() npu_backend torchair.get_npu_backend() model torch.compile(model, backendnpu_backend, dynamicFalse) model(cache, x, slot_mapping) print(cache.shape, cache.dtype)源码实现剖析算子定义与属性注册算子核心定义在 op_host/kv_compress_epilog_def.cpp输入cacheUINT8/ND、xBF16/ND、slot_mappingINT32/INT64/ND输出cache与输入同 tensor 语义原地更新。InferShape 与 InferDtype图模式下 GE 依赖 op_host/kv_compress_epilog_infershape.cpp 推导 NetOutput shape输出 cache 的 shape 与 dtype 直接拷贝输入 cache缺省空实现会导致 GE NetOutput shape 为 [] 与 FX 图不一致单算子直调路径不经过该函数。Tiling 策略tiling 实现op_host/kv_compress_epilog_tiling_arch35.cpp的核心逻辑平台信息获取读取 AIV 核心数与 UB 内存大小输入 shape 校验bs由 slotMapping 各维累乘得到d取 x 末维校验d 64、d ≤ 8192、d % 64 0行负载分配rowOfFormerBlock CeilDiv(bs, coreNum)将 bs 行均分到各核心UB 容量适配通过二分查找同文件 L316-L331计算每个核心单次循环可搬入的行数 rowFactor考量 x 输入缓冲double buffer、cache 输出缓冲、tmp 缓冲与 mode2 的 scratch 缓冲之和不超过 ubSize分页 cache 布局支持GetCacheViewLayout解析非连续 blockNum 维的 strideblockStride ≥ blockSize × headDim 校验支持各 block 不紧密排布的分页场景后处理设置 block 维度使用的核心数、单一 tiling keymode0/1/2 在 kernel 内按 quantMode 运行时分支与 workspace 大小并输出 tiling 数据。Kernel 实现kernel 入口op_kernel/kv_compress_epilog.cpp通过GET_TILING_DATA_WITH_STRUCT解析 tiling 数据按 tiling key 分发到KvCompressEpilogOps::KvCompressEpilog的Init/Process流程并保护性地保存/恢复浮点溢出模式控制寄存器FLOAT_OVERFLOW_MODE_CTRL。具体的量化路径per-group amax 计算、FP8/FP4 转换、roundScale 取 2 的幂、hifloat8 静态量化等在 op_kernel/kv_compress_epilog_kernel.h 与 op_kernel/kv_compress_epilog_common.h 中实现。测试与验证仓库为算子提供了完整的单元测试Host 侧单测tests/ut/op_host/test_kv_compress_epilog_infershape.cpp 验证 InferShape/InferDtype 行为tests/ut/op_host/test_kv_compress_epilog_tiling.cpp 验证 tiling 数据bs/d/kvCacheCol/rowFactor 等的生成正确性Kernel 侧单测tests/ut/op_kernel/arch35/test_kv_compress_epilog.cpp 在 arch35Ascend 950上对三种量化模式做数值级验证工程样例examples/test_aclnn_kv_compress_epilog.cppaclnn 单算子与 examples/test_geir_kv_compress_epilog.cppGE 图模式。完整的算子级说明含调用方式矩阵可参见模块 READMEattention/kv_compress_epilog/README.md。该算子默认支持确定性计算相关规范可参考 确定性计算。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考