Fast KVzip：基于门控KV淘汰的高效准确LLM推理

用轻量门控模块直接从隐藏状态预测KV重要性，免去运行时重建开销

前置知识

KV缓存（Key-Value Cache）

在Transformer自回归推理中，每生成一个新token都需要对之前所有token做注意力计算。为避免重复计算，系统会把每一层注意力的Key和Value向量缓存下来，这就是KV缓存。它使得生成第n个token时只需计算当前token的Query与所有缓存Key的点积，而不用重新处理整个序列。KV缓存的大小与序列长度和模型层数成正比，是长上下文推理的主要内存瓶颈。

本文的核心目标就是压缩KV缓存，理解其结构和内存占用机制是理解全文的基础。

注意力稀疏性（Attention Sparsity）

研究发现，LLM在做注意力计算时，并非所有token都同等重要。对于给定的query，通常只有少数token会获得较高的注意力权重，其余token的贡献微乎其及。这种稀疏性在不同层和不同头之间表现不同——有些头专注于局部信息，有些头关注全局摘要。这种现象是KV缓存压缩的理论基础：如果能识别出哪些KV对不重要，就可以安全地淘汰它们。

Fast KVzip正是利用这种稀疏性，通过门控机制预测哪些KV对会被未来使用，哪些可以安全丢弃。

注意力汇聚（Attention Sink）

Xiao等人（2024）发现，即使初始token（如BOS token）在语义上不重要，LLM也会给它们分配异常高的注意力权重。这些token被称为attention sinks。这是因为softmax归一化需要某些token吸收多余的注意力概率质量。移除这些token会导致模型性能显著下降。这一发现启发了许多保留初始token的KV压缩策略。

本文的门控架构设计借鉴了attention sink的思想，引入可学习的sink key来建模KV重要性。

分块预填充（Chunked Prefill）

处理长上下文时，一次性计算所有token的KV缓存会消耗大量GPU内存。分块预填充将输入序列切成固定大小的块（如16K tokens），逐块计算KV特征并存入缓存。这样可以控制峰值内存使用量，使长上下文处理在有限显存的GPU上成为可能。Agrawal等人（2024）的Sarathi-Serve系统展示了这种技术在推理服务中的实用性。

Fast KVzip在每个预填充块之后执行KV淘汰，与分块预填充无缝集成，这是其内存效率的关键。

KVzip的重建方法

KVzip（Kim等人，2025a）是一种query-agnostic的KV缓存压缩方法。它通过让模型重建上下文来评估每个KV对的重要性：给定一段文本，模型用KV缓存重建该文本，记录每个KV对在重建过程中接收到的最大注意力分数，以此作为重要性指标。这种方法泛化性好，但代价是推理时需要额外的前向传播，相当于将预填充计算量翻倍。

Fast KVzip的核心贡献就是用轻量门控替代KVzip的运行时重建过程，将重建开销转移到离线训练阶段。

研究动机

随着LLM上下文窗口不断扩展到数十万甚至百万token级别，KV缓存的线性内存增长成为实际部署的核心瓶颈。现有的KV缓存压缩方法面临一个两难困境：计算开销可以忽略的方法（如H2O、SnapKV等基于启发式规则的方法）会导致显著的性能下降，因为它们依赖推理时观察到的注意力模式来决定淘汰哪些KV对，而这些模式往往过拟合于当前输入，在多查询场景下泛化能力差。另一方面，能保持准确性的方法（如KVzip）需要在推理时执行重建过程，这使得预填充阶段的计算量翻倍，对于延迟敏感的部署场景来说代价过高。具体来说，KVzip在处理170K token的上下文时，预填充时间约为无压缩基线的两倍，这在实时交互场景中是不可接受的。因此，如何在不牺牲准确性的前提下实现快速KV压缩，是一个亟待解决的实际问题。

本文的目标是本文的目标是设计一种KV缓存压缩方法，使其同时满足三个条件：第一，在KV缓存预算低至原始大小的25%-30%时仍能保持接近无压缩的模型性能；第二，压缩过程的计算和内存开销可以忽略不计，不增加预填充延迟；第三，方法具有通用性，能跨不同模型架构（Qwen、Gemma等）、不同任务类型（检索、问答、代码理解、数学推理）和不同推理阶段（预填充和解码）保持一致的效果。作者希望将KVzip的准确性优势与其运行时效率问题解耦，实现两全其美。

与已有工作不同的是，本文的独特视角在于一个关键洞察：KV对的未来利用程度本质上是一个内在属性，可以直接从输入隐藏状态中解码出来，而不需要像KVzip那样重建整个上下文来评估。这个洞察意味着，我们可以在模型的前向传播过程中直接插入一个轻量级的门控模块来预测KV重要性，而不是在推理时额外执行一次完整的重建过程。作者进一步发现，现有的多种KV压缩方法（如DuoAttention、Expected Attention、FastGen）都可以被重新表述为广义门控框架的受限实例，这表明门控机制提供了一种统一且原则性的方法来理解和改进KV压缩，将问题从启发式设计转变为数据驱动优化。

核心方法

Fast KVzip的核心思路可以用一个类比来理解：想象你是一个图书馆管理员，需要在有限的书架空间里保留最有价值的书籍。传统方法（KVzip）的做法是每次有读者来借书时，你都把所有书翻一遍来判断哪些书重要——准确但太慢。Fast KVzip的做法则是训练一个经验丰富的助手（门控模块），这个助手只需看书的封面信息（隐藏状态）就能快速判断这本书未来的使用概率，然后直接做出保留或淘汰的决定。这个助手的判断能力来自于对大量书籍使用记录的学习，而不是每次都重新评估。技术路线上，Fast KVzip在每个Transformer注意力层中插入一个轻量级的门控分支，该分支接收当前层的隐藏状态作为输入，输出每个KV头的重要性分数（0到1之间），然后根据这些分数淘汰低重要性的KV对。门控模块使用KVzip的重建分数作为训练目标进行离线训练，训练完成后推理时只需一次前向传播即可完成压缩。

Fast KVzip与已有方法最本质的区别在于压缩决策的时间和方式。KVzip在推理时通过重建整个上下文来评估KV重要性，这相当于做了一次额外的前向传播；而Fast KVzip将这个评估过程转移到离线训练阶段，训练一个轻量门控模块来学会如何从隐藏状态中直接预测KV重要性。这个转变的关键依据是：KV对的未来利用程度是输入的内在属性，可以被模型的隐藏状态所编码。具体来说，门控模块采用一种受attention sink启发的低秩汇聚注意力架构：给定隐藏状态 $h \in \mathbb{R}^D$，通过线性投影得到查询向量 $q \in \mathbb{R}^{G \times H \times D'}$ 和键向量 $k \in \mathbb{R}^{H \times D'}$，然后通过类似softmax的机制计算重要性分数。这个设计使得门控模块可以利用注意力机制的表达能力来建模复杂的KV重要性模式，同时保持极低的参数量（每层仅约几千个参数）和计算开销。

方法步骤详情

Fast KVzip的完整工作流程分为训练阶段和推理阶段。训练阶段：首先从FineWeb-Edu预训练语料中随机采样10K到30K token长度的序列，构建总计1M token的训练集。然后，对每个训练样本，使用KVzip的重建方法计算每个KV对在每一层接收到的最大注意力分数，作为训练目标。接着，在每一层独立训练门控模块，使用二元交叉熵损失和随机梯度下降，学习率为0.2，训练5000步，batch size为1000。LLM权重在整个过程中保持冻结，只训练门控参数。推理阶段：在预填充时，采用16K token的分块大小，每个分块计算完KV特征后，门控模块同步计算重要性分数并淘汰低分KV对，同时保留最近4K token的局部窗口以利用注意力的局部性偏差。在解码时，维护一个128 token的隐藏状态缓冲区，定期（缓冲区满时）并行计算门控分数并执行KV淘汰，保留最近128 token的局部窗口。这种缓冲区策略将门控开销降低到模型前向传播延迟的约1%。

技术新颖性

Fast KVzip的技术新颖性体现在三个方面。第一，它首次系统地证明了KV对的未来利用程度可以从输入隐藏状态中直接预测，而不需要重建过程，这一洞察从根本上改变了KV压缩的计算范式。第二，它提出了一种低秩汇聚注意力门控架构，借鉴了attention sink的机制，使用可学习的sink key来建模KV重要性，这种架构在检索任务上显著优于简单的线性模型或MLP。第三，它提供了一个统一的理论框架来理解现有KV压缩方法：DuoAttention可以看作输入无关的常数门控，Expected Attention可以看作基于键状态的二次函数门控，FastGen可以看作基于输入token的布尔门控。这种统一视角将KV压缩从启发式设计转化为数据驱动优化，为未来的方法设计提供了原则性指导。此外，与Lorcret（仅支持预填充压缩）和TrimKV（在预填充场景下性能下降）不同，Fast KVzip在预填充和解码两个阶段都实现了有效压缩。

实验结果

Fast KVzip在多个维度上展现了出色的性能。在预填充密集型任务上，以Qwen2.5-7B-1M为测试模型，在SCBench的12个数据集上，Fast KVzip在KV缓存预算为30%-40%时就能保持接近完整缓存的性能，与KVzip（需要两倍预填充计算量）持平，并显著优于SnapKV、DuoAttention、Expected Attention等基线。在RULER-4K基准测试中，Fast KVzip在25%的KV预算下仍能匹配KVzip的压缩性能。在KVPress基准测试（2025年12月版本）上，Fast KVzip的表现超过了当时所有SOTA方法。在解码密集型任务上，以Qwen3-8B和Qwen3-14B为测试模型，在AIME24和MATH数学推理基准上，Fast KVzip在KV预算为4096 token时实现了接近无损的性能，而Qwen3的early-stopping-of-thinking策略会导致推理性能急剧下降。在效率方面，对于170K token的上下文，Fast KVzip相比无压缩基线减少了峰值KV内存使用（从36GB降至约12GB），同时预填充时间从KVzip的约120秒降至约60秒，与无压缩基线相当。在模型通用性上，在Qwen3-8B、Qwen3-14B、Qwen2.5-14B-1M、Gemma3-12B和Qwen3-8B-FP8（量化模型）上都验证了有效性。门控训练仅需0.59-0.83 H100 GPU小时，门控参数存储仅需0.11-0.30 GB。

查看结构化数据

任务	指标	本文	基线	提升
SCBench检索任务（Retr.KV）	准确率	约68%（30% KV预算，Qwen2.5-7B-1M）	KVzip约67%，SnapKV约40%，DuoAttention约25%	匹配KVzip，显著优于其他基线
SQuAD问答	准确率	约92%（30% KV预算，Qwen2.5-7B-1M）	KVzip约92%，SnapKV约82%，DuoAttention约78%	匹配KVzip，比SnapKV高约10个百分点
AIME24数学推理	平均分	约70%（4K KV预算，Qwen3-8B）	R-KV约55%，Early Stopping约20%	比R-KV高约15个百分点
RULER-4K	准确率	约88%（25% KV预算）	KVzip约88%，KVPress约78%	匹配KVzip，优于KVPress约10个百分点
MATH数学推理	准确率	约88%（4K KV预算，Qwen3-8B）	R-KV约82%，Early Stopping约80%	比R-KV高约6个百分点
预填充效率（170K tokens）	预填充时间（秒）	约60秒	KVzip约120秒，无压缩约60秒	比KVzip快2倍，与无压缩持平

局限与改进

尽管Fast KVzip展现了出色的性能，但仍存在一些局限性。首先，门控模块的训练需要使用KVzip的重建过程来生成目标分数，这意味着训练阶段仍然依赖于这个计算密集型的过程，尽管这只是一次性的离线成本。其次，作者在论文中承认，将门控机制应用于预训练阶段（而非仅在推理时使用）是一个重要的未来方向，当前方法仅适用于冻结权重的模型，无法利用门控来诱导硬件高效的结构化计算模式。第三，论文中提到的并发工作KVzap探索了类似的门控优化目标，虽然具体架构和训练配方不同，但两者在同一时期独立发展，说明这是一个被广泛关注的研究方向。从实验设计来看，论文主要在英文数据集上进行评估，对多语言场景的适用性尚未验证。此外，门控模块虽然轻量，但对于超大规模模型（如数百B参数），每层数千参数的门控模块累积起来的存储和计算开销是否仍然可以忽略，还需要进一步验证。

独立分析的弱点

从独立分析的角度来看，Fast KVzip存在几个值得关注的弱点。首先，门控模块使用二元交叉熵损失训练，这意味着它将KV重要性预测视为一个二分类问题（重要/不重要），但实际上重要性是一个连续值。这种离散化可能会损失信息，尤其是在边界情况下。一个改进方向是使用回归损失（如MSE或Huber损失）来直接预测连续的重要性分数，或者使用排序损失来更好地保持相对重要性顺序。其次，当前方法使用固定的局部窗口大小（预填充4K，解码128），但在不同任务和不同层中，最优的局部窗口大小可能不同。可以考虑让门控模块自适应地预测每个层的最优窗口大小。第三，门控模块在每个层独立训练，没有考虑层间信息的传递。引入跨层的门控信息共享可能会提高预测准确性，尤其是在深层的复杂注意力模式上。第四，对于混合注意力机制的模型（如Gemma3-12B的滑动窗口注意力），Fast KVzip只压缩全局注意力的KV缓存，这限制了其在这些模型上的压缩潜力。

未来方向

作者提出了几个重要的未来方向。第一，将门控机制应用于预训练阶段，让模型在训练时就学会产生硬件高效的结构化计算模式，这可能会带来比推理时压缩更大的效率提升。第二，探索一个统一框架，将Fast KVzip的门控机制与其他门控方法（如MoD的计算跳过、MoBA的稀疏注意力）集成，支持多种决策：跳过计算、丢弃KV特征、或选择性存储KV以供后续检索。第三，扩展门控机制使其支持多选择决策，而不是简单的保留/淘汰二元选择。基于当前成果还可以延伸出更多方向：将门控机制与KV量化结合，实现更极致的压缩比；探索在线学习的门控策略，让门控模块在推理过程中持续适应当前上下文的特点；将方法扩展到多模态LLM，处理图像、视频等模态的KV缓存压缩。

复现评估

从复现角度来看，Fast KVzip的复现条件相当友好。代码已在GitHub开源（https://github.com/Janghyun1230/fastkvzip），这大大降低了复现门槛。门控训练数据使用公开的FineWeb-Edu数据集，采样1M token即可，不需要特殊的数据准备。训练资源需求极低：对于14B参数规模的模型，门控训练仅需不到1 H100 GPU小时，门控参数存储仅需0.11-0.30 GB。推理时使用标准的PyTorch和FlashAttention-2，不需要特殊的硬件或软件环境。论文评估了多种开源模型（Qwen2.5、Qwen3、Gemma3），这些模型都可以公开获取。超参数设置统一（S=16, D'=16），不需要针对不同模型进行调优。整体而言，一个有LLM推理经验的研究团队应该能在1-2天内完成复现。