企业官网建设流程全解析

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 系列专属特性，当前仅支持以下产品，其余产品均不支持：

该结论与 op_host/kv_compress_epilog_def.cpp 中AICore().AddConfig("ascend950")的注册配置一致，即算子仅在 ascend950 架构上编译与运行。

功能说明

算子语义

在 KV Cache 的 Epilog 阶段，对 KV Cache 进行原地压缩更新：

• 将bfloat16激活值x量化压缩后，按slotMapping散写到 cache；
• slotMapping中值为-1的 token 跳过不处理；
• x尾轴d的后 64 列为 rope 段、前d-64列为 nope 段；
• cache 为原地更新（in-place），未被命中的行保留原值。

支持三种量化模式：

1. quantMode=0（group，bf16 scale）：逐组动态量化压缩为 FP8，scale存储为bfloat16，rope 段保留bfloat16；
2. quantMode=1（group，e8m0 scale）：量化方法同模式 0，但scale存储为float8_e8m0（2 的幂指数，即 MX-FP8 微缩放语义）；roundScale=true时scale向上取到 2 的幂；
3. quantMode=2（rope hifloat8 静态 + nope FLOAT4_E2M1 动态）：rope 段（后 64 列）做 hifloat8静态量化（乘xScale），nope 段（前d-64列）做 per-groupFLOAT4_E2M1 动态量化（FP4_MAX=6.0），scale以bfloat16写出，roundScale在该模式下不生效。

计算公式

对x的最后一维（d 轴）按组计算每组的 amax 并量化为目标 dtype，记第 g 组为 $x_g$：

场景 1（quantMode=0）：group(64) 量化为 FP8，scale存储为bfloat16；rope 段保留bfloat16。

$$ scale_g = \frac{\max(|x_g|)}{FP8_MAX}, \quad q_i = \mathrm{round}\left(\frac{x_i}{scale_g}\right) $$

场景 2（quantMode=1）：同场景 1，但scale存储为float8_e8m0，roundScale=true时scale向上取到 2 的幂。

场景 3（quantMode=2）：rope 段（后 64 列）做 hifloat8 静态量化，nope 段（前 d-64 列）做 per-group FLOAT4_E2M1 动态量化（FP4_MAX=6.0），scale以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 量化默认采用 E4M3FN（max=448.0），同时保留了 E5M2（max=57344.0）路径的常量定义。

典型示例（shape 速查）

cache shape: [128, 16, 1, 384] # 4D [blockNum, blockSize, 1, headDim], num_slots=2048, 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 参数说明

参数明细

属性默认值的源码依据

各属性的默认值在 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 返回码。第一段接口完成入参校验，出现以下场景时报错：

对应实现见 op_host/op_api/aclnn_kv_compress_epilog.cpp：先做OP_CHECK_NULL空指针检查，再做IsEmpty()空 Tensor 检查。

aclnnKvCompressEpilog 参数说明

返回值同样为 aclnnStatus 状态码（参见 aclnn 返回码）。

约束说明

• 确定性计算：aclnnKvCompressEpilog 默认为确定性实现。
• cache 仅支持四维 shape[blockNum, blockSize, 1, headDim]（num_slots = blockNum × blockSize），倒数第二维固定为 1，不支持其他维数；仅在 blockNum 维支持非连续（分页场景下各 block 可不紧密排布，但 block 内部须连续）。
• headDim 约束：cache 末维 headDim 须 ≥ 每行写出字节数 kvCacheCol = 对齐(concatCol)：
  • concatCol 计算（quantMode=0/1）：(d-64) + 128 + ⌈(d-64)/64⌉ × scaleBytes，其中 scaleBytes 在 mode0=2（bf16 scale）、mode1=1（e8m0 scale）；
  • concatCol 计算（quantMode=2）：64 + (d-64)/2 + ((d-64)/quantGroupSize) × 2；
  • 对齐规则：quantMode=1 不补齐（kvCacheCol = concatCol）；quantMode=0/2 按 32B 对齐（kvCacheCol = ⌈concatCol/32⌉ × 32）。
  • 示例：d=256、quantMode=1 → kvCacheCol=323 → headDim ≥ 323；d=512、quantMode=1 → kvCacheCol=583 → headDim ≥ 583。
• slotMapping 的维度应等于 x 的维度减 1，即 slotMapping 为 x 除最后一维外的所有维度展平。
• x 的最后一维（d 轴）需满足 d % 64 == 0 且 64 < d ≤ 8192，按每 64 个元素一组进行逐组量化。
• quantMode=2 时，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.cpp）：

if (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_cast<int64_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，每组附带 scale；mode2 下 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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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::vector<int64_t> cacheShape = {128, 16, 1, 384}; // 4D KV Cache [blockNum,blockSize,1,headDim]，num_slots=2048 std::vector<int64_t> xShape = {1024, 256}; // 待量化激活 std::vector<int64_t> slotShape = {1024}; // slot 索引映射 void* cacheDeviceAddr = nullptr; void* xDeviceAddr = nullptr; void* slotDeviceAddr = nullptr; aclTensor* cache = nullptr; aclTensor* x = nullptr; aclTensor* slotMapping = nullptr; std::vector<uint8_t> cacheHostData(GetShapeSize(cacheShape), 0); std::vector<uint16_t> xHostData(GetShapeSize(xShape), 0); // bfloat16 以 uint16 承载 std::vector<int32_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.py），API 详见 torchapi_kv_compress_epilog.md：

cann_ops_transformer.kv_compress_epilog(cache, x, slot_mapping, *, quant_group_size=64, quant_mode="fp8_e8m0", round_scale=True, x_scale=1.0) -> None

要点：

• 原地语义：该接口无返回值，结果直接写回输入cache（schema 中标注Tensor(a!)），调用后直接使用入参 cache；
• quant_mode 为字符串（大小写不敏感），内部经KvCompressQuantMode枚举映射为算子侧 int："fp8_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须 >= kvCacheCol（d=256/quant_mode=1 时kvCacheCol=323，取head_dim=384 满足） block_num, block_size, head_dim = 128, 16, 384 bs, d = 1024, 256 cache = torch.zeros(block_num, block_size, 1, head_dim, dtype=torch.uint8).npu() x = torch.randn(bs, d, dtype=torch.bfloat16).npu() slot_mapping = torch.randint(0, block_num * block_size, (bs,), dtype=torch.int32).npu() kv_compress_epilog( cache, x, slot_mapping, quant_group_size=64, quant_mode="fp8_e8m0", round_scale=True, x_scale=1.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_size=64, quant_mode="fp8_e8m0", round_scale=True, x_scale=1.0) return cache model = KvCompressEpilogModel().npu() npu_backend = torchair.get_npu_backend() model = torch.compile(model, backend=npu_backend, dynamic=False) model(cache, x, slot_mapping) print(cache.shape, cache.dtype)

源码实现剖析

算子定义与属性注册

算子核心定义在 op_host/kv_compress_epilog_def.cpp：输入cache（UINT8/ND）、x（BF16/ND）、slot_mapping（INT32/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）的核心逻辑：

1. 平台信息获取：读取 AIV 核心数与 UB 内存大小；
2. 输入 shape 校验：bs由 slotMapping 各维累乘得到，d取 x 末维，校验d > 64、d ≤ 8192、d % 64 == 0；
3. 行负载分配：rowOfFormerBlock = CeilDiv(bs, coreNum)将 bs 行均分到各核心；
4. UB 容量适配：通过二分查找（同文件 L316-L331）计算每个核心单次循环可搬入的行数 rowFactor，考量 x 输入缓冲（double buffer）、cache 输出缓冲、tmp 缓冲与 mode2 的 scratch 缓冲之和不超过 ubSize；
5. 分页 cache 布局支持：GetCacheViewLayout解析非连续 blockNum 维的 stride（blockStride ≥ blockSize × headDim 校验），支持各 block 不紧密排布的分页场景；
6. 后处理：设置 block 维度（使用的核心数）、单一 tiling key（mode0/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 在 arch35（Ascend 950）上对三种量化模式做数值级验证；
• 工程样例：examples/test_aclnn_kv_compress_epilog.cpp（aclnn 单算子）与 examples/test_geir_kv_compress_epilog.cpp（GE 图模式）。

完整的算子级说明（含调用方式矩阵）可参见模块 README：attention/kv_compress_epilog/README.md。该算子默认支持确定性计算，相关规范可参考 确定性计算。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer