由于Transformer模型中的键值（KV）缓存呈线性增长，部署长上下文LLM的成本非常高。例如，使用Llama 3.1-70B以float16格式处理100万 tokens 需要高达330GB的内存。kvpress实现了多种KV缓存压缩方法，并基于🤗 Transformers进行了基准测试，旨在为该领域的研究人员和开发者简化新方法的开发流程。

安装

pip install kvpress

若需本地安装，请使用 uv：

git clone https://github.com/NVIDIA/kvpress.git
cd kvpress
uv sync

如需安装所有可选依赖项，请执行以下命令：

git clone https://github.com/NVIDIA/kvpress.git
cd kvpress
uv sync --extra eval --extra flash-attn

使用

KVPress提供了一系列在预填充阶段压缩KV缓存的“压缩器”。每个压缩器都关联有一个compression_ratio属性，用于衡量缓存的压缩程度。使用压缩器最简单的方式是通过我们自定义的KVPressTextGenerationPipeline。当导入kvpress时，它会自动注册为名为“kv-press-text-generation”的Transformers流水线，并为您处理聊天模板和分词：

from transformers import pipeline
from kvpress import ExpectedAttentionPress

model = "Qwen/Qwen3-8B"
pipe = pipeline("kv-press-text-generation", model=model, device_map="auto", dtype="auto")

context = "一段您希望一次性压缩的超长文本"
question = "\n关于压缩后上下文的一个问题"  # 可选

press = ExpectedAttentionPress(compression_ratio=0.5)
answer = pipe(context, question=question, press=press)["answer"]

在上述代码片段中，压缩仅应用于上下文tokens，以便您可以针对不同问题评估压缩效果。更多详细示例请参阅Wikipedia笔记本演示，该演示也在Colab上提供此处。

from transformers import pipeline
from kvpress import KnormPress
from kvpress import DecodingPress

# 初始化流水线
device = "cuda:0"
model = "meta-llama/Llama-3.1-8B-Instruct"
model_kwargs = {"attn_implementation": "flash_attention_2"}
pipe = pipeline("kv-press-text-generation", model=model，device=device，model_kwargs=model_kwargs)

# 创建一个每10步压缩一次，目标缓存大小为512 tokens 的解码压缩器
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_steps=10，
    token_buffer_size=512
)

# 与流水线结合使用
context = "一段您希望在生成过程中压缩的超长文本"
question = "请根据这段上下文讲一个长故事"
response = pipe(context，question=question，press=decoding_press)["answer"]

可用的压缩器

当前所有的压缩器都不需要训练，并且都继承自 BasePress（源代码）。

其中，若干压缩器继承自 ScorerPress（源代码），它们依赖于一个分数来裁剪重要性最低的 KV 对：

• RandomPress（源代码）：随机分数
• KnormPress（源代码，论文）：键的逆范数
• SnapKVPress（源代码，论文）：最后查询的平均注意力权重
• ExpectedAttentionPress（源代码，笔记本）：生成阶段的期望注意力权重
• StreamingLLMPress（源代码，论文）：仅保留初始和最近的标记
• TOVAPress（源代码，论文）：对所有头的最后查询注意力权重取平均
• ObservedAttentionPress（源代码，论文）：在预填充阶段观察到的平均注意力权重
• QFilterPress（源代码，论文）：将键表示投影到查询向量的主要奇异值分解成分上，以近似计算注意力分数。
• PyramidKVPress（源代码，论文）：维持金字塔形的缓存大小，为较低层分配更多缓存预算，而为较高层分配较少。
• LagKVPress（源代码，论文）：利用 KV 延迟相关信息进行压缩。它无需查询、无需注意力权重，且兼容 Flash Attention。
• KeyDiffPress（源代码，论文）：仅根据键的相似性驱逐标记。
• NonCausalAttnPress（源代码，论文）：基于非因果分块注意力分数驱逐标记。
• LeverageScorePress（源代码，论文）：根据近似的统计杠杆度驱逐标记（即保留键空间中的异常值）。
• CompactorPress（源代码，论文）：根据压缩比混合使用 NonCausalAttnPress 和 LeverageScorePress。
• CURPress（源代码，论文）：基于 CUR 分解和近似杠杆度分数修剪键和值。
• KVzapPress（源代码，论文，训练）：通过快速代理模型近似 KVzip+。需与 DMSPress 配合使用。
• FastKVzipPress（源代码，论文）：通过轻量级门控机制近似 KVzip。

还有一些压缩器采用不同的逻辑：

• ThinKPress（源代码，论文）：根据最后查询上的通道注意力分数压缩键的维度。
• SimLayerKVPress（源代码，论文）：识别“懒惰”层，并对其应用 StreamingLLM 方法。
• DuoAttentionPress（源代码，论文）：将注意力头分为检索头（不压缩）和流式头（采用 StreamingLLM 方法）。
• FinchPress（源代码，论文）：类似于 SnapKV，但具有动态窗口大小和键值重旋转。
• KVzipPress（源代码，论文）：通过上下文重建识别冗余的 KV 对。以多次前向传播为代价，实现近乎无损的压缩。

最后，我们提供了一些可以与其他压缩器组合使用的包装型压缩器：

• AdaKVPress（源代码，论文）：裁剪任何 ScorerPress 的底部分数，但跨所有头执行，从而实现按头的压缩。
• PerLayerCompressionPress（源代码）：为每一层使用不同的压缩比进行压缩（实验性功能）。
• ComposedPress（源代码）：通过串联多个压缩器的前向钩子将其组合在一起。
• KeyRerotationPress（源代码）：对已压缩的键进行重新旋转，以保持连续的 RoPE 嵌入。
• ChunkKVPress（源代码，论文）：通过选择重要区块进行压缩，同时保持语义连贯性。
• ChunkPress（源代码，论文）：分别对每个序列区块的 KV 缓存进行压缩。这可以在长序列中实现更均匀的压缩效果。
• CriticalKVPress 和 CriticalAdaKVPress（源代码，论文）：利用 Wo @ 值的 L1 范数细化分数，并结合两阶段选择机制。
• BlockPress（源代码，论文）：将输入序列划分为不重叠的块，并迭代式地进行压缩（⚠️ 并非真正的分块预填充实现）。
• DecodingPress（源代码）：允许在解码过程中进行压缩，请参阅本 README 中的解码部分。
• PrefillDecodingPress（源代码）：允许在预填充和解码过程中同时进行压缩。
• DMSPress（源代码，论文）：对于任何 ScorerPress，驱逐分数低于给定阈值的键和值，而不是依赖于 top-k 分数。支持预填充和解码（若 decoding=True），但仅支持密集预填充，不支持稀疏预填充。

有关现有 KV 缓存压缩方法的详细列表，请参阅 Awesome-KV-Cache-Compression 或 Awesome-LLM-Compression。

评估

我们提供了一个简单的命令行界面，用于在多个长上下文数据集上评估不同压缩方法的性能。

• 准确性：使用我们的 CLI 直接在流行的基准测试上测试您的方法。
• 速度与内存：speed_and_memory 笔记本可以帮助您测量峰值内存使用量和总时间节省。

有关更多详细信息和结果，请参阅此仓库中的 evaluation 目录。

以下是我们从 报告的不同压缩方法在 RULER 数据集中、4k 上下文长度下的平均性能。

量化

我们通过 transformers 的 QuantizedCache 类支持 KV 缓存量化（参见 HF 博客文章）。要使用它，只需将缓存对象传递给您的管道：

from transformers import QuantizedCache

cache = QuantizedCache(backend="quanto", nbits=4)

pipe(..., cache=cache)

默认情况下，使用的是 DynamicCache（不进行量化）。

[!重要]
要使用 QuantizedCache, 您需要安装额外的依赖项（例如 pip install optimum-quanto）。

贡献

我们欢迎贡献！要添加新的压缩方法，只需打开一个问题或提交拉取请求。请查看 new_press.ipynb 笔记本，获取分步指南。

引用

如果您在研究中使用 KVPress，请引用我们的论文：

@article{devoto2025expectedattention,
  title={Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution},
  author={Devoto, Alessio and Jeblick, Maximilian and J{\'e}gou, Simon},
  journal={arXiv preprint arXiv:2510.00636},
  year={2025},
  url={https://arxiv.org/abs/2510.00636}
}

常见问题解答

pipe = pipeline("kv-press-text-generation", model=model, device_map="auto")

import torch
from transformers import AutoModelForCausalLM
from kvpress import KnormPress

device = "cuda:0"
ckpt = "meta-llama/Meta-Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(ckpt).to(device)
press = KnormPress(compression_ratio=0.4)

inputs = model.dummy_inputs["input_ids"].to(device)

with torch.no_grad():
    print(model(inputs).past_key_values[0][0].shape)
    # torch.Size([3, 8, 5, 128])
    
with torch.no_grad(), press(model):
    print(model(inputs).past_key_values[0][0].shape)
    # torch.Size([3, 8, 3, 128])

with press(model):
    outputs = model.generate(inputs)

from transformers import pipeline
from kvpress import KnormPress
from kvpress import DecodingPress

# 初始化管道
device = "cuda:0"
model = "meta-llama/Llama-3.1-8B-Instruct"
model_kwargs = {"attn_implementation": "flash_attention_2"}
pipe = pipeline("kv-press-text-generation", model=model, device=device, model_kwargs=model_kwargs)

# 创建一个每 10 步压缩到 512 个 token 的解码压缩方法
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_steps=10,
    token_buffer_size=512
)

# 与管道一起使用
context = "一段非常长的文本，您希望在生成过程中对其进行压缩"
question = "请根据这段上下文讲一个长故事"
response = pipe(context, question=question, press=decoding_press)["answer"]

from transformers import pipeline
from kvpress import CriticalKVPress、KnormPress
from kvpress import DecodingPress、PrefillDecodingPress

# 初始化管道
device = "cuda:0"
model = "meta-llama/Llama-3.1-8B-Instruct"
model_kwargs = {"attn_implementation": "flash_attention_2"}
pipe = pipeline("kv-press-text-generation", model=model, device=device, model_kwargs=model_kwargs)

# 预填充和解码采用不同的策略
prefill_press = CriticalKVPress(KnormPress())
decoding_press = DecodingPress(
    base_press=KnormPress(compression_ratio=0.2),
    compression_steps=5,
    token_buffer_size=256
)

# 将两者结合
combined_press = PrefillDecodingPress(
    prefilling_press=prefill_press,
    decoding_press=decoding_press
)

context = "一段非常长的上下文，将在预填充阶段被压缩"
question = "生成一份详细的分析报告，该报告将在解码阶段被压缩"
response = pipe(context，question=question，press=combined_press)["answer"]

kvpress 快速上手指南

kvpress 是一个基于 🤗 transformers 的开源工具库，旨在通过多种压缩方法减少大语言模型（LLM）在长上下文场景下的 KV 缓存显存占用，从而降低部署成本并提升推理效率。

环境准备

• 系统要求：Linux 或 macOS 环境，推荐配备 NVIDIA GPU 以加速推理。
• Python 版本：建议 Python 3.9 及以上。
• 核心依赖：
  • transformers (Hugging Face)
  • torch (PyTorch)
  • accelerate
• 可选依赖：若需极致性能，建议安装 flash-attn (Flash Attention 2)。

国内加速建议：在安装依赖时，推荐使用清华或阿里镜像源以提升下载速度。

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

方式一：通过 PyPI 安装（推荐）

最简便的安装方式，适用于大多数用户：

pip install kvpress

方式二：源码安装（含可选依赖）

如果你需要本地开发或使用评估工具及 Flash Attention，建议使用 uv 进行源码安装：

# 克隆仓库
git clone https://github.com/NVIDIA/kvpress.git
cd kvpress

# 同步基础环境 (需先安装 uv: pip install uv)
uv sync

# 同步包含评估工具和 flash-attn 的完整环境
uv sync --extra eval --extra flash-attn

基本使用

kvpress 的核心概念是"Press"（压缩器），它可以在预填充（prefilling）阶段压缩 KV 缓存。最简单的方式是使用封装好的 KVPressTextGenerationPipeline，它自动处理分词和对话模板。

以下示例演示如何使用 ExpectedAttentionPress 将上下文压缩 50%：

from transformers import pipeline
from kvpress import ExpectedAttentionPress

# 1. 初始化管道
# 管道名称固定为 "kv-press-text-generation"
model = "Qwen/Qwen3-8B"
pipe = pipeline("kv-press-text-generation", model=model, device_map="auto", dtype="auto")

# 2. 定义长上下文和问题
context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"  # 可选，用于测试不同问题下的压缩效果

# 3. 配置压缩策略
# compression_ratio=0.5 表示保留 50% 的 KV 对
press = ExpectedAttentionPress(compression_ratio=0.5)

# 4. 执行推理
# 压缩仅应用于 context 部分，question 部分保持完整以确保回答质量
answer = pipe(context, question=question, press=press)["answer"]

print(answer)

进阶：解码阶段压缩（实验性功能）

除了预填充阶段，kvpress 还支持在生成（decoding）过程中定期压缩缓存。这通过 DecodingPress 包装器实现：

from transformers import pipeline
from kvpress import KnormPress, DecodingPress

# 初始化管道 (建议使用 flash_attention_2)
model = "meta-llama/Llama-3.1-8B-Instruct"
pipe = pipeline(
    "kv-press-text-generation", 
    model=model, 
    device="cuda:0", 
    model_kwargs={"attn_implementation": "flash_attention_2"}
)

# 创建解码压缩器
# 每生成 10 个 token 压缩一次，目标是将缓存大小维持在 512 个 token
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_steps=10,
    token_buffer_size=512
)

context = "A very long text you want to compress during generation"
question = "Tell me a long story about this context"

response = pipe(context, question=question, press=decoding_press)["answer"]

注意：并非所有压缩算法都支持解码阶段压缩，目前主要支持继承自 ScorerPress 的算法（如 KnormPress, CriticalKVPress 等）。