随着上下文变长，模型需要保存的KV Cache也随之增加。为降低缓存开销，DeepSeek-V4.1通过CSA2中的Compressor，将连续多个token的特征汇聚成一个压缩表示，供后续注意力计算使用。这一过程中，Compressor算子既要处理不同长度的输入，也要保存尚未完成压缩的数据，供下一次调用继续计算。
借助CANNBot的triton-op-generator插件，通过给Agent提供开源PyTorch标杆并描述特定要求，完成了Compressor的 Triton-Ascend代码生成、精度验证和性能优化，得到了可以用于实际业务的算子。
在Device侧性能测试中，128K Prefill预填充阶段的耗时从 24.357 ms 降至 7.341 ms，加速约3.32倍；32路Decode解码阶段在产生压缩输出时，耗时从 150.66 μs 降至 23.03 μs，加速约6.54倍。
Compressor的计算流程是对输入做两路投影：一路提取特征kv，另一路计算门控分数score。每凑齐r个token，就根据分数计算权重，对特征加权汇聚，再经过RMSNorm归一化，得到一个压缩表示。
图 1：Compressor 主路径（r > 1）。r = 1 时仅保留特征投影与归一化。
流式调用需要保存未完成的压缩组。例如 r = 2 时，第一个 token 到来后，先保存投影结果；第二个token到来后，再将两者汇聚，产生压缩输出。
实现时还要处理以下情况：
输入规模不同。Prefill 一次处理大量 token；Decode 每个请求通常只有一个 token，两者需要不同的任务安排。
状态要跨调用保留。 一次计算结束后，尚未凑齐一组的数据要留给下次使用。
请求位置会变化。 请求在batch中换了位置，也要通过映射表找到原来的状态。
序列长度不一定相同。变长打包输入需要结合 cu_seqlens、seqused判断每条序列的起点和有效长度。
这些情况需要在接口定义、精度验证、性能采集和问题定位中一并考虑。
triton-op-generator是CANNBot面向Triton-Ascend算子生成与优化的官方插件。插件按阶段组织任务构建、算法设计、代码生成、验证、优化和报告归档，由编排器调用相应的Skill执行。
它支持三种常见输入方式：
提供 PyTorch 标杆，生成 Triton-Ascend 实现。
同时提供 PyTorch 标杆和 GPU Triton kernel，分别用于精度验证和算法结构参考。
提供 GPU Triton kernel 和输入数据，将已有 GPU 实现迁移到 Triton-Ascend。
本文采用第一种方式，在昇腾A2/A3系列产品上进行Compressor算子的开发与优化。
triton-op-generator支持Claude Code、OpenCode、Codex、Trae 等交互工具。以Claude Code为例，可通过统一安装器安装：
npx @cannbot-plugin/cannbot@latest install triton-op-generator --tool claude也可以从源码安装：
git clone --recurse-submodules --shallow-submodules https://gitcode.com/cann/cannbot.gitcd cannbot/plugins-official/triton-op-generatorbash init.sh project claude
安装过程会同步准备 Triton 系列 Skills、CANNBot 知识库和工作流提示词，无需再逐个安装依赖 Skill。
运行前还需安装配套的CANN Toolkit、PyTorch、torch_npu和 Triton-Ascend环境，并确保目标昇腾硬件与当前环境一致。
在项目根目录启动Agent交互工具，提供PyTorch 标杆、目标接口和优化目标，即可启动任务。本次 Compressor 任务可以这样描述：
使用 CANNBot 的 triton-op-generator，以 DeepSeek-V4.1 model.py 中的Compressor PyTorch 实现为标杆，生成 Triton-Ascend 算子。目标接口需要支持kv_state、score_state、state_block_table、cu_seqlens 和 seqused；覆盖 Prefill 与 Decode 场景。先完成精度对齐，再以相对 PyTorch 路径的加速比为指标进行性能优化。
为精确理解需求，防止返工，提示词中推荐详细说明接口语义。例如，state_block_table = [2, 0, 1]表示三条请求分别使用状态池第 2、0、1 行。请求在 batch 中换位后，仍按映射表读取原来的状态；只写“支持状态缓存”，无法明确这种寻址规则。
任务启动后，Agent会编排工作流推进各个阶段，并将算子实现、report.md、summary.json和会话记录写入独立任务目录。
生成工作流分为Phase 0—7 八个阶段。下图展示各阶段的分工，以及代码生成和性能优化中的迭代关系。
图 2：triton-op-generator完整工作流
Phase 0确认目标硬件、提供哪些参考代码，以及按什么指标优化。本次Compressor生成任务以PyTorch实现为标杆，目标是生成Triton-Ascend实现，并优化Device侧执行耗时。
Phase 1由triton-task-extractor整理接口和测试用例。Compressor算子除了输入、输出，还要明确历史状态如何保留、未凑齐的压缩组如何继续计算，以及请求顺序变化后如何找到对应状态。未明确的规则单独记录，供后续核查。
各方面都确定后工作流会固定 PyTorch 标杆，并在每次验证和性能采集前核对文件哈希，确保每次比较基准保持一致。
triton-op-designer会先参考知识库中同类算子的约束和问题，构建算法草图，阐明各步计算如何衔接、数据从哪里读取，以及如何分配给多个计算核并行执行。
就Compressor算子而言，不同输入布局需要先转换成统一的寻址信息，再分别为 Prefill 和 Decode 设计不同的计算调度。Prefill 输入数据量大，重点减少投影和池化的数据搬运；Decode 输入数据量小，重点增加可并行的任务、减少启动开销。其次针对没有凑齐压缩组的场景，只更新状态，不执行池化。
Agent综合上述设计形成算法草图，供下一阶段生成代码使用。
triton-op-coding根据算法草图生成代码，生成后先检查是否存在 PyTorch 完成计算等作弊行为。发现这类问题后，返回修复。
通过检查后，triton-op-verifier使用相同输入与初始状态，将Triton-Ascend结果与PyTorch标杆进行精度比较。除常规输入外，验证场景还覆盖变长序列、请求换位、未完成压缩组的续算，以及多次连续调用等情况。
Compressor算子验证过程中曾发现缓存并发写入乱序、负索引引发越界读取等问题。工作流会要求Agent根据失败用例定位并修复算子实现，再重新验证。直到全部用例通过后，测量当前实现的性能，作为 Phase 4 的比较基线。
进入性能优化阶段后，triton-latency-optimizer负责查找优化点并修改实现代码。每轮修改后，由triton-op-verifier重新验证精度，再测量性能。每轮性能结果与此前的最佳版本比较。性能改善的方案予以保留，没有改善的改动则回退并记录原因。
需要进一步定位性能瓶颈时，由triton-simulator-optimizer采集并分析指令执行情况；优化过程中也会查看IR（编译器中间表示），确认代码经过编译后产生了哪些搬运和同步操作。诊断结果用于指导代码修改优化方向。
本次优化主要涉及Prefill投影、Decode投影和池化三个环节：
图 3：三个环节的瓶颈、改动和实测结果。
Prefill：减少频繁搬运。指令流水诊断显示，初版投影的大量时间花在细粒度装载上。在片上容量允许的范围内增大分块、减少搬运次数后，独立投影耗时从3.870 ms降至 3.059 ms。
Decode：提高并行度。Decode场景沿用大矩阵配置时，计算任务只分配给少了量核，其余核未能参与。指令流水诊断发现，循环中的搬运与同步等待开销较高。针对这两个问题，优化器沿输出列拆分计算，让更多核参与，同时增大每轮处理的输入通道块，减少循环次数。双路投影耗时从 28.3 μs降至14.8 μs。
池化：避免不必要的读取与同步。指令流水诊断结合IR分析发现，历史状态读取和门控分数的重复装载带来了额外开销。从两方面优化：将需要接续历史数据的首组单独处理，让后续完整组直接使用本次输入，并让一次计算连续处理多个压缩组；针对 r = 2 的完整组，两个token的门控分数各读取一次，供后续计算复用，将装载操作从8次减至4次，减少相应的同步开销。经过多轮调整，128K场景的池化耗时从约6.2 ms降至1.065ms。
算子优化结束后，Phase 5对选定的最终代码重新执行完整精度验证，输出代码和报告。报告中的性能数据对应同一份代码、同一组测试用例，并说明是否达到优化目标。
随后Phase 6导出开发会话，保留修改、回退和未采用方案的记录，便于复查设计取舍。
Phase 7在任务成功完成且知识库可用时，将经过验证的实现要求、算法方案和性能数据整理到项目内知识库，供后续同类算子开发参考。
测试在昇腾A2系列产品上进行，PyTorch标杆和Triton-Ascend实现采用相同权重、输入和初始状态，每组压缩的token数 r = 2。PyTorch参考路径的投影使用 FP32，Triton-Ascend实现使用BF16乘法与FP32累加。
下表统计Device侧kernel耗时之和。加速比反映了数据类型、并行方式及计算划分等调整的综合收益。
Decode数据统计的是产生压缩输出时的单步耗时。B表示批大小，例如B = 32 对应32条请求共同执行这一步的总耗时。只更新状态、不产生压缩输出的步骤耗时未计入。
以上为所列用例的算子侧结果。Python调用和启动间隙仍会影响端到端耗时，模型服务中的吞吐与时延收益还需进一步验证。
这次使用triton-op-generator开发Compressor，交付内容包括算子代码，以及对应的接口说明、验证记录和优化依据。Prefill、Decode和池化各自根据测量结果调整，每轮改动都经过精度与性能复验，开发者可以据此了解设计取舍和结果适用范围。
DeepSeek-V4.1官方Compressor源码：门控汇聚、流式状态与归一化实现
https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash/blob/517ef62/inference/model.py#L429
CANNBot源码与插件入口：
https://gitcode.com/cann/cannbot
