CANN PyPTO 算子实战Compressor 长序列 KV Cache 压缩算子的原理、参数与调用方法【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-inferCompressor 是 CANN 推理优化样例库cann-recipes-infer中面向长序列推理场景的关键融合算子其核心思路是把每 4 个或 128 个 token 的 KV cache 压缩为 1 个再让后续 token 与这些压缩后的 KV cache 执行 DSADynamic Sparse Attention计算从而显著削减长序列下的显存占用与计算开销。本文基于 ops/pypto_python/docs/Compressor.md 展开结合 ops/pypto_python/impl/compressor_pypto.py 的 kernel 实现与 ops/pypto_python/example/test_compressor_pypto.py 的测试样例系统讲解该算子的产品支持情况、数学流程、函数原型、全部参数约束、三种 ratio/rotate 组合的 kernel 分支以及如何在 DeepSeek-V4 整网中启用 PyPTO 路径。产品支持情况Compressor 算子基于昇腾硬件平台提供支持的产品型号如下产品是否支持Atlas A2 推理系列产品√Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT 推理系列产品√对应的环境要求可参考 ops/pypto_python/README.mdAtlas A2/A3 系列使用 Linux ARM 操作系统镜像版本为cann9.0.pt2.8.0_ds_pypto_aarch_image:v0.2驱动版本 25.5.0可用npu-smi info确认Ascend 950PR/DT 系列的环境与驱动待发布。功能说明与核心计算流程API 功能定位功能Compressor 将每 4 或 128 个 token 的 KV cache 压缩成一个随后每个 token 与这些压缩后的 KV cache 进行 DSA 计算。在长序列场景下Compressor 可以有效减少计算开销同时降低 KV cache 存储开销。归属从 ops/pypto_python/impl/compressor_pypto.py 的模块注释可见该实现面向 deepseek_v4 模型设计属于 DeepSeek-V4 整网中 IndexerLILightning Indexer与 SFASparse Flash Attention两类路径共用的压缩前处理算子。主要计算过程算子内部计算分为三步投影与打分将输入 $X$ 与 $W^{KV}$ 做 Matmul 得到 $kv_state$将输入 $X$ 与 $W^{Gate}$ 做 Matmul 后再与 $Ape$positional biases做 Add 得到 $score_state$$kv_state$ 与 $score_state$ 根据输入的start_pos完成更新写入 cache。加权聚合对 $kv_state$ 和 $score_state$ 进行数据重排再对 $score_state$ 做 softmax将 softmax 结果与 $kv_state$ 做 Mul随后进行 Reducesum得到压缩后的 KV 表示。后处理根据输入数据norm_weight、rope_sin、rope_cos进行 RmsNorm 和 ROPE 运算并根据rotate决定是否需要额外进行 Hadamard Transform最终得到 $cmp_kv$ 结果输出。函数原型Compressor 通过torch.library注册为pypto命名空间下的 FRAGMENT 算子见 impl/compressor_pypto.py对外暴露的原型如下torch.ops.pypto.compressor( x, kv_state, score_state, kv_block_table, state_block_table, sin, cos, wkv, wgate, ape, weight, hadamard, start_pos, ratio, rope_head_dim, rotate ) - (Tensor)Python 侧封装入口为compressor_pypto()其内部调用torch.ops.pypto.compressor(...)对应 impl/compressor_pypto.py并同时注册了Meta与NPU两个实现NPU实现经allow_in_graph暴露的npu_compressor()进入 PyPTO kernel 派发impl/compressor_pypto.py。参数说明下表逐项说明每个参数的含义、数据类型、shape 与约束源自原文档并结合check_args校验逻辑佐证参数必选/可选含义数据类型Shapex必选原始不经压缩的数据对应公式中的 $X$。不支持非连续数据格式支持NDbfloat16[B, S, H]kv_state必选kv_state的历史数据对应公式中的 $kv_state$。不支持非连续NDfloat32[block_num, block_size, coff*D]score_state必选score_state的历史数据对应公式中的 $score_state$。不支持非连续NDfloat32[block_num, block_size, coff*D]kv_block_table必选kv_state历史数据的 page table。不支持非连续NDint32[B, ceil(max_S/block_size)]state_block_table必选score_state历史数据的 page table。不支持非连续NDint32[B, ceil(max_S/block_size)]sin必选Rope 计算的权重系数bfloat16[min(T, T//ratioB), rope_head_dim]cos必选Rope 计算的权重系数bfloat16[min(T, T//ratioB), rope_head_dim]wkv必选KV 和压缩权重的权重参数对应公式中的 $W^{KV}$。不支持非连续NDbfloat16[coff*D, H]wgate必选KV 和压缩权重的权重参数对应公式中的 $W^{Gate}$。不支持非连续NDbfloat16[coff*D, H]ape必选输入的 positional biases对应公式中的 $Ape$。不支持非连续NDfloat32[ratio, coff*D]weight必选计算 RmsNorm 时的权重系数bfloat16[D,]hadamard可选Hadamard Transform 的权重矩阵。不支持非连续NDbfloat16[D, D]start_pos可选计算起始位置。不支持非连续ND。输入为 None 时表示从 0 开始计算int32[B,]ratio必选数据压缩率int支持 4 / 128rope_head_dim必选rope_cos和rope_sin的 hidden 层最小单元int目前仅支持 64rotate必选是否需要额外进行 Hadamard Transformbool—其中coff 1 (ratio 4)即 ratio4 时为 2引入 overlap 通道见下文ratio128 时为 1。check_args还进一步校验了 dtypex/cos/sin/hadamard/weight/wkv/wgate必须为bfloat16kv_state/score_state/ape必须为float32kv_block_table/score_block_table/start_pos必须为int32。返回值说明outTensor必选输出表示压缩后的数据。不支持非连续数据格式支持ND数据类型支持bfloat16输出 shape 为[min(T, T // ratio B), D]。不压缩的条目的输出数据值为零——即尚未凑满一个压缩窗口的请求其输出位置以 0 填充对应 kernel 中pypto.assemble(zero, ...)与torch.zeros初始化逻辑见 impl/compressor_pypto.py。约束说明该接口支持 B 泛化batch 维度动态。S 支持 1/2/3/4即单次调用最多处理 4 个 token与 ratio4 或 128 的压缩窗口配合check_args中校验x.size(1) in [1, 2, 3, 4]。D 支持 128 / 512与rotate联动d 128 and rotate或d 512 and not rotate。H 支持 4096。block_size支持 128。ratio仅支持 4 或 128rope_head_dim仅支持 64rotate必须为 bool见 impl/compressor_pypto.py 的断言。调用方法与源码实现解读运行示例在已配置 PyPTO 环境镜像cann9.0.pt2.8.0_ds_pypto_aarch_image:v0.2的容器中执行python3 ops/pypto_python/example/test_compressor_pypto.py该测试脚本覆盖三个用例见 example/test_compressor_pypto.pytest_comp_128ratio128、rotateFalsebsz64、start_pos8192、seq2、d512模拟长序列 decode 压缩test_comp_4ratio4、rotateFalsebsz64、start_pos2、seq2、d512test_comp_indexerratio4、rotateTruebsz64、start_pos1、seq2、d128对应 Indexer 路径。每个用例都会调用golden_compress()构造 CPU/参考实现并用assert_allclose分别对kv_state、score_statertol/atol1e-3以及压缩输出rtol≈7.8e-3、atol1e-4做数值比对。kernel 派发逻辑npu_compressor()依据rotate与ratio的组合派发到三个 PyPTO JIT kernelimpl/compressor_pypto.py组合调用的 kernel特点rotateTrue且ratio4compressor_ratio_4_rotate_kerneld128压缩完成后额外执行out_view hadamardHadamard Transform输出为 bf16rotateFalse且ratio4compressor_ratio_4_kerneld512coff2含 overlap 通道rotateFalse且ratio128compressor_ratio_128_kerneld512coff1三个 kernel 均使用pypto.jit编译runtime_options中开启了device_sched_mode: 3与triple_stream_sched: True以配合多流调度。kernel 内部按 batch 分块b64双层循环外层LOOP_COMP_1遍历 batch 块内层LOOP_COMP_2遍历块内每条请求。压缩窗口与状态更新ratio4 为例以compressor_ratio_4_kernel为例内部关键流程为将输入xreshape 为[bsz*s1, h]与wkv、wgate做 Matmulb_transTrue输出 FP32得到kv_t、score_timpl/compressor_pypto.py。读取每条请求的start_pos判断start_pos % ratio s1 ratio无需压缩将每个 token 的kv/score加上对应位置的ape通过pypto.assemble直接写入kv_state_out/score_state_out需要压缩先按 page table 取当前窗口kv_state/score_state当start_pos ratio时用 0 和-inf填充前序窗口保证首窗口 softmax 正确否则取上一窗口的[pre_kv_block_idx, pre_start]数据随后逐个 token 执行scatter_update_3d更新窗口内状态当pos ratio - 1时把前序窗口与当前窗口的 d 维后半段拼接成[b, 8, d]ratio4、coff2 的 overlap 结构计算kv_state_tmp * softmax(score_state_tmp, 1)并沿维度 1 求和得到压缩后的[b, d]impl/compressor_pypto.py。对压缩结果依次执行rms_normepsilon1e-6FP32 计算后转回原 dtype、拆出kv_nope与kv_rope对kv_rope做 interleaved 二维 RoPEinterleaved_rope_2d含rotate_half与transpose/reshape变换最后拼接并assemble到输出impl/compressor_pypto.py。ratio128 的 kernel 流程一致但 coff1直接在单 block128 槽位内完成 128 个 token 的状态聚合不再拼接前序窗口impl/compressor_pypto.py。算子级辅助函数kernel 依赖三个可复用的 PyPTO 函数softmax(x, dim)以amax数值稳定版实现减最大值、exp、sum、divrms_norm(input, gamma, epsilon1e-6)FP32 计算均方根归一化后再乘 gammarotate_half/interleaved_rope_2d实现 interleave 模式的 RoPE 旋转。DeepSeek-V4 整网集成方式Compressor 算子已支持集成到 DeepSeek-V4 整网ops/pypto_python/README.md 的DeepSeek-V4 整网集成样例执行章节启用步骤如下权重与数据集准备参考 models/deepseek_v4/README.md 中的模型权重和数据集准备章节。代码修改适配参考 models/deepseek_v4/README.md 进行网络配置适配。修改 yaml 配置在models/deepseek_v4/config/路径下需要执行的 yaml 文件中修改model_configmodel_config: enable_limit_core: False enable_pypto: True其中enable_pypto: True将 PyPTO 算子路径集成到网络enable_limit_core需关闭。需要注意二者互斥enable_limit_coreTrue时要求enable_multi_streamsTrue且不支持enable_pyptoTrue见 models/deepseek_v4/README.md同时enable_pypto目前仅支持platform_version a3models/deepseek_v4/models/modeling_deepseek.py。拉起推理参考 models/deepseek_v4/README.md 的拉起多卡推理章节执行结束出现model run success即表示推理成功。在整网层面Compressor 模块位于 models/deepseek_v4/models/modules/compressor.pyCompressor类持有wkv/wgate两个复制的线性层、ape位置偏置参数、normRmsNorm与hadamard_matrix并通过compressor_prolog压缩计算 RmsNorm partial RoPE与compressor_epilog将压缩结果写回 LI/SFA 压缩 cache支持 float8/hifloat8 量化 cache组织整网调用forward中还在 CPcontext parallel场景下通过 all-gather 与状态块搬运完成跨 rank 的压缩状态同步。数值精度与验证要点测试脚本example/test_compressor_pypto.py提供了可直接对照的参考实现golden_compress其核心语义与 kernel 一致压缩触发条件为(start_pos i 1) % ratio 0即窗口内最后一个 token 触发压缩ratio4 时采用 overlap 窗口把前序窗口的[:d]与当前窗口的[d:]拼接成 8 行再 softmax 加权对应 kernel 中pre_kv_state/pre_score_state的逻辑ratio128 时直接对当前 block 内 128 行 softmax 加权压缩后统一执行 RmsNorm并对尾部rope_head_dim维度做 interleave RoPErotateTrue时再乘 Hadamard 矩阵。通过比对kv_state、score_state与压缩输出的误差可以验证算子端到端的数值正确性。这也是在实际业务中接入该算子时建议保留的验证流程。小结Compressor 算子以每 4/128 个 token 压缩为一个 KV cache为骨架在单个 PyPTO kernel 内融合了 Matmul、状态读写page table assemble/scatter、softmax 加权聚合、RmsNorm、RoPE 与可选的 Hadamard Transform是长序列 DeepSeek-V4 推理降低 KV cache 开销与计算量的关键一环。开发者可按 ops/pypto_python/docs/Compressor.md 的参数约束直接调用torch.ops.pypto.compressor或通过enable_pypto: True在 DeepSeek-V4 整网中一键启用。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN PyPTO 算子实战：Compressor 长序列 KV Cache 压缩算子的原理、参数与调用方法
CANN PyPTO 算子实战Compressor 长序列 KV Cache 压缩算子的原理、参数与调用方法【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-inferCompressor 是 CANN 推理优化样例库cann-recipes-infer中面向长序列推理场景的关键融合算子其核心思路是把每 4 个或 128 个 token 的 KV cache 压缩为 1 个再让后续 token 与这些压缩后的 KV cache 执行 DSADynamic Sparse Attention计算从而显著削减长序列下的显存占用与计算开销。本文基于 ops/pypto_python/docs/Compressor.md 展开结合 ops/pypto_python/impl/compressor_pypto.py 的 kernel 实现与 ops/pypto_python/example/test_compressor_pypto.py 的测试样例系统讲解该算子的产品支持情况、数学流程、函数原型、全部参数约束、三种 ratio/rotate 组合的 kernel 分支以及如何在 DeepSeek-V4 整网中启用 PyPTO 路径。产品支持情况Compressor 算子基于昇腾硬件平台提供支持的产品型号如下产品是否支持Atlas A2 推理系列产品√Atlas A3 推理系列产品√Ascend 950PR / Ascend 950DT 推理系列产品√对应的环境要求可参考 ops/pypto_python/README.mdAtlas A2/A3 系列使用 Linux ARM 操作系统镜像版本为cann9.0.pt2.8.0_ds_pypto_aarch_image:v0.2驱动版本 25.5.0可用npu-smi info确认Ascend 950PR/DT 系列的环境与驱动待发布。功能说明与核心计算流程API 功能定位功能Compressor 将每 4 或 128 个 token 的 KV cache 压缩成一个随后每个 token 与这些压缩后的 KV cache 进行 DSA 计算。在长序列场景下Compressor 可以有效减少计算开销同时降低 KV cache 存储开销。归属从 ops/pypto_python/impl/compressor_pypto.py 的模块注释可见该实现面向 deepseek_v4 模型设计属于 DeepSeek-V4 整网中 IndexerLILightning Indexer与 SFASparse Flash Attention两类路径共用的压缩前处理算子。主要计算过程算子内部计算分为三步投影与打分将输入 $X$ 与 $W^{KV}$ 做 Matmul 得到 $kv_state$将输入 $X$ 与 $W^{Gate}$ 做 Matmul 后再与 $Ape$positional biases做 Add 得到 $score_state$$kv_state$ 与 $score_state$ 根据输入的start_pos完成更新写入 cache。加权聚合对 $kv_state$ 和 $score_state$ 进行数据重排再对 $score_state$ 做 softmax将 softmax 结果与 $kv_state$ 做 Mul随后进行 Reducesum得到压缩后的 KV 表示。后处理根据输入数据norm_weight、rope_sin、rope_cos进行 RmsNorm 和 ROPE 运算并根据rotate决定是否需要额外进行 Hadamard Transform最终得到 $cmp_kv$ 结果输出。函数原型Compressor 通过torch.library注册为pypto命名空间下的 FRAGMENT 算子见 impl/compressor_pypto.py对外暴露的原型如下torch.ops.pypto.compressor( x, kv_state, score_state, kv_block_table, state_block_table, sin, cos, wkv, wgate, ape, weight, hadamard, start_pos, ratio, rope_head_dim, rotate ) - (Tensor)Python 侧封装入口为compressor_pypto()其内部调用torch.ops.pypto.compressor(...)对应 impl/compressor_pypto.py并同时注册了Meta与NPU两个实现NPU实现经allow_in_graph暴露的npu_compressor()进入 PyPTO kernel 派发impl/compressor_pypto.py。参数说明下表逐项说明每个参数的含义、数据类型、shape 与约束源自原文档并结合check_args校验逻辑佐证参数必选/可选含义数据类型Shapex必选原始不经压缩的数据对应公式中的 $X$。不支持非连续数据格式支持NDbfloat16[B, S, H]kv_state必选kv_state的历史数据对应公式中的 $kv_state$。不支持非连续NDfloat32[block_num, block_size, coff*D]score_state必选score_state的历史数据对应公式中的 $score_state$。不支持非连续NDfloat32[block_num, block_size, coff*D]kv_block_table必选kv_state历史数据的 page table。不支持非连续NDint32[B, ceil(max_S/block_size)]state_block_table必选score_state历史数据的 page table。不支持非连续NDint32[B, ceil(max_S/block_size)]sin必选Rope 计算的权重系数bfloat16[min(T, T//ratioB), rope_head_dim]cos必选Rope 计算的权重系数bfloat16[min(T, T//ratioB), rope_head_dim]wkv必选KV 和压缩权重的权重参数对应公式中的 $W^{KV}$。不支持非连续NDbfloat16[coff*D, H]wgate必选KV 和压缩权重的权重参数对应公式中的 $W^{Gate}$。不支持非连续NDbfloat16[coff*D, H]ape必选输入的 positional biases对应公式中的 $Ape$。不支持非连续NDfloat32[ratio, coff*D]weight必选计算 RmsNorm 时的权重系数bfloat16[D,]hadamard可选Hadamard Transform 的权重矩阵。不支持非连续NDbfloat16[D, D]start_pos可选计算起始位置。不支持非连续ND。输入为 None 时表示从 0 开始计算int32[B,]ratio必选数据压缩率int支持 4 / 128rope_head_dim必选rope_cos和rope_sin的 hidden 层最小单元int目前仅支持 64rotate必选是否需要额外进行 Hadamard Transformbool—其中coff 1 (ratio 4)即 ratio4 时为 2引入 overlap 通道见下文ratio128 时为 1。check_args还进一步校验了 dtypex/cos/sin/hadamard/weight/wkv/wgate必须为bfloat16kv_state/score_state/ape必须为float32kv_block_table/score_block_table/start_pos必须为int32。返回值说明outTensor必选输出表示压缩后的数据。不支持非连续数据格式支持ND数据类型支持bfloat16输出 shape 为[min(T, T // ratio B), D]。不压缩的条目的输出数据值为零——即尚未凑满一个压缩窗口的请求其输出位置以 0 填充对应 kernel 中pypto.assemble(zero, ...)与torch.zeros初始化逻辑见 impl/compressor_pypto.py。约束说明该接口支持 B 泛化batch 维度动态。S 支持 1/2/3/4即单次调用最多处理 4 个 token与 ratio4 或 128 的压缩窗口配合check_args中校验x.size(1) in [1, 2, 3, 4]。D 支持 128 / 512与rotate联动d 128 and rotate或d 512 and not rotate。H 支持 4096。block_size支持 128。ratio仅支持 4 或 128rope_head_dim仅支持 64rotate必须为 bool见 impl/compressor_pypto.py 的断言。调用方法与源码实现解读运行示例在已配置 PyPTO 环境镜像cann9.0.pt2.8.0_ds_pypto_aarch_image:v0.2的容器中执行python3 ops/pypto_python/example/test_compressor_pypto.py该测试脚本覆盖三个用例见 example/test_compressor_pypto.pytest_comp_128ratio128、rotateFalsebsz64、start_pos8192、seq2、d512模拟长序列 decode 压缩test_comp_4ratio4、rotateFalsebsz64、start_pos2、seq2、d512test_comp_indexerratio4、rotateTruebsz64、start_pos1、seq2、d128对应 Indexer 路径。每个用例都会调用golden_compress()构造 CPU/参考实现并用assert_allclose分别对kv_state、score_statertol/atol1e-3以及压缩输出rtol≈7.8e-3、atol1e-4做数值比对。kernel 派发逻辑npu_compressor()依据rotate与ratio的组合派发到三个 PyPTO JIT kernelimpl/compressor_pypto.py组合调用的 kernel特点rotateTrue且ratio4compressor_ratio_4_rotate_kerneld128压缩完成后额外执行out_view hadamardHadamard Transform输出为 bf16rotateFalse且ratio4compressor_ratio_4_kerneld512coff2含 overlap 通道rotateFalse且ratio128compressor_ratio_128_kerneld512coff1三个 kernel 均使用pypto.jit编译runtime_options中开启了device_sched_mode: 3与triple_stream_sched: True以配合多流调度。kernel 内部按 batch 分块b64双层循环外层LOOP_COMP_1遍历 batch 块内层LOOP_COMP_2遍历块内每条请求。压缩窗口与状态更新ratio4 为例以compressor_ratio_4_kernel为例内部关键流程为将输入xreshape 为[bsz*s1, h]与wkv、wgate做 Matmulb_transTrue输出 FP32得到kv_t、score_timpl/compressor_pypto.py。读取每条请求的start_pos判断start_pos % ratio s1 ratio无需压缩将每个 token 的kv/score加上对应位置的ape通过pypto.assemble直接写入kv_state_out/score_state_out需要压缩先按 page table 取当前窗口kv_state/score_state当start_pos ratio时用 0 和-inf填充前序窗口保证首窗口 softmax 正确否则取上一窗口的[pre_kv_block_idx, pre_start]数据随后逐个 token 执行scatter_update_3d更新窗口内状态当pos ratio - 1时把前序窗口与当前窗口的 d 维后半段拼接成[b, 8, d]ratio4、coff2 的 overlap 结构计算kv_state_tmp * softmax(score_state_tmp, 1)并沿维度 1 求和得到压缩后的[b, d]impl/compressor_pypto.py。对压缩结果依次执行rms_normepsilon1e-6FP32 计算后转回原 dtype、拆出kv_nope与kv_rope对kv_rope做 interleaved 二维 RoPEinterleaved_rope_2d含rotate_half与transpose/reshape变换最后拼接并assemble到输出impl/compressor_pypto.py。ratio128 的 kernel 流程一致但 coff1直接在单 block128 槽位内完成 128 个 token 的状态聚合不再拼接前序窗口impl/compressor_pypto.py。算子级辅助函数kernel 依赖三个可复用的 PyPTO 函数softmax(x, dim)以amax数值稳定版实现减最大值、exp、sum、divrms_norm(input, gamma, epsilon1e-6)FP32 计算均方根归一化后再乘 gammarotate_half/interleaved_rope_2d实现 interleave 模式的 RoPE 旋转。DeepSeek-V4 整网集成方式Compressor 算子已支持集成到 DeepSeek-V4 整网ops/pypto_python/README.md 的DeepSeek-V4 整网集成样例执行章节启用步骤如下权重与数据集准备参考 models/deepseek_v4/README.md 中的模型权重和数据集准备章节。代码修改适配参考 models/deepseek_v4/README.md 进行网络配置适配。修改 yaml 配置在models/deepseek_v4/config/路径下需要执行的 yaml 文件中修改model_configmodel_config: enable_limit_core: False enable_pypto: True其中enable_pypto: True将 PyPTO 算子路径集成到网络enable_limit_core需关闭。需要注意二者互斥enable_limit_coreTrue时要求enable_multi_streamsTrue且不支持enable_pyptoTrue见 models/deepseek_v4/README.md同时enable_pypto目前仅支持platform_version a3models/deepseek_v4/models/modeling_deepseek.py。拉起推理参考 models/deepseek_v4/README.md 的拉起多卡推理章节执行结束出现model run success即表示推理成功。在整网层面Compressor 模块位于 models/deepseek_v4/models/modules/compressor.pyCompressor类持有wkv/wgate两个复制的线性层、ape位置偏置参数、normRmsNorm与hadamard_matrix并通过compressor_prolog压缩计算 RmsNorm partial RoPE与compressor_epilog将压缩结果写回 LI/SFA 压缩 cache支持 float8/hifloat8 量化 cache组织整网调用forward中还在 CPcontext parallel场景下通过 all-gather 与状态块搬运完成跨 rank 的压缩状态同步。数值精度与验证要点测试脚本example/test_compressor_pypto.py提供了可直接对照的参考实现golden_compress其核心语义与 kernel 一致压缩触发条件为(start_pos i 1) % ratio 0即窗口内最后一个 token 触发压缩ratio4 时采用 overlap 窗口把前序窗口的[:d]与当前窗口的[d:]拼接成 8 行再 softmax 加权对应 kernel 中pre_kv_state/pre_score_state的逻辑ratio128 时直接对当前 block 内 128 行 softmax 加权压缩后统一执行 RmsNorm并对尾部rope_head_dim维度做 interleave RoPErotateTrue时再乘 Hadamard 矩阵。通过比对kv_state、score_state与压缩输出的误差可以验证算子端到端的数值正确性。这也是在实际业务中接入该算子时建议保留的验证流程。小结Compressor 算子以每 4/128 个 token 压缩为一个 KV cache为骨架在单个 PyPTO kernel 内融合了 Matmul、状态读写page table assemble/scatter、softmax 加权聚合、RmsNorm、RoPE 与可选的 Hadamard Transform是长序列 DeepSeek-V4 推理降低 KV cache 开销与计算量的关键一环。开发者可按 ops/pypto_python/docs/Compressor.md 的参数约束直接调用torch.ops.pypto.compressor或通过enable_pypto: True在 DeepSeek-V4 整网中一键启用。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考