NVIDIA/kvpress

KVPress – 长上下文LLM的KV缓存压缩

它是什么

• 一个Python库，实现了多种按压（press）算法，这些算法在预填充阶段（以及可选地在解码阶段）对基于Transformer的语言模型的键值（KV）缓存进行剪枝或压缩。减少KV缓存可以降低内存使用，并可能加快超长上下文（例如，100万token）的推理速度。

为什么重要

• 在现代LLM中，KV缓存随处理token数量线性增长。在100万token的文档上运行70B模型可能需要超过300GB的GPU内存。KVPress让研究人员和开发人员无需重写模型代码即可尝试多种最先进的压缩方法。

主要特性

• 即插即用管道 – KVPressTextGenerationPipeline 注册为Hugging Face管道（kv-press-text-generation）。
• 丰富的按压目录 – 提供超过30种即用方法，从简单的随机剪枝到复杂的基于注意力和矩阵分解的技术（例如，KnormPress、SnapKVPress、KVzipPress、CURPress）。
• 预填充和解码压缩 – DecodingPress 及相关包装器允许在生成token时进行周期性压缩。
• 可组合 – 包装按压（AdaKVPress、ComposedPress、ChunkPress 等）可让您堆叠或组合策略。
• 无需训练 – 所列的所有按压均无需额外微调模型即可工作。
• 量化缓存支持 – 与Hugging Face的QuantizedCache（例如，通过optimum-quanto实现4位量化）配合使用。
• 多GPU就绪 – 与accelerate设备映射兼容。
• 基准测试与排行榜 – 提供评估CLI和笔记本，以及RULER 4k-token基准上的公开HF Space排行榜。

安装

pip install kvpress               # 基本安装
# 或者，用于开发，包含可选依赖（评估、flash-attention）
uv sync --extra eval --extra flash-attn

基本使用示例

from transformers import pipeline
from kvpress import ExpectedAttentionPress

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

context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"
press = ExpectedAttentionPress(compression_ratio=0.5)

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

该按压在模型关注问题之前压缩上下文token的KV缓存，使您能够比较不同的压缩比例。

评估

• 提供一个CLI，可在长上下文数据集上运行每个按压，并报告准确性、内存减少和加速情况。
• 仓库包含用于交互式性能分析的笔记本（speed_and_memory.ipynb、wikipedia_demo.ipynb）。
• RULER 4k-token基准的结果显示在README中链接的Hugging Face排行榜上。

文档与资源

• Colab笔记本演示、HF Space演示和博客文章介绍了该库。
• 每个按压链接到其源文件和相应的arXiv论文，以便深入学习。
• 一个“新按压”笔记本指导贡献者添加自定义压缩方法。

许可证

• Apache 2.0（参见README中的徽章）。

引用

@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}
}

获取方式

• GitHub: https://github.com/NVIDIA/kvpress
• PyPI: kvpress
• Hugging Face Space: https://huggingface.co/spaces/nvidia/kvpress