kvpress

环境要求

• Python 环境（可用 pip 安装包）
• 本地源码安装需先安装 uv（uv sync 使用）
• 运行压缩推理需要 transformers 与 PyTorch
• 长上下文模型推理需要足够显存（README 举例 1M tokens 的 Llama 3.1-70B float16 需 330GB）
• 多 GPU 推理需 accelerate（README 说明通过 accelerate 支持）

安装与启动步骤

1. 1安装 kvpress

   最简安装方式，直接用 pip 从 PyPI 安装 kvpress 包。

   pip install kvpress

   复制

2. 2本地源码安装

   使用 uv 从 GitHub 克隆仓库并同步依赖，适合本地开发与修改源码。

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

   复制

3. 3安装可选依赖

   需要评估工具或 flash-attn 时，用 --extra 安装可选依赖，普通使用可跳过。

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

   复制

4. 4运行压缩推理

   导入 kvpress 后 pipeline 自动注册为 kv-press-text-generation，用 press 对上下文做压缩。

   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 = "A very long text you want to compress once and for all"
   question = "
   A question about the compressed context"
   
   press = ExpectedAttentionPress(compression_ratio=0.5)
   answer = pipe(context, question=question, press=press)["answer"]

   复制

5. 5启用解码阶段压缩

   实验特性 DecodingPress 包装 ScorerPress，在生成过程中周期性压缩 KV 缓存。

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

   复制

6. 6多 GPU 推理

   README 说明 kvpress 通过 accelerate 支持多 GPU，设置 device_map="auto" 即可。

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

   复制

关键配置

如何确认成功

运行示例代码不报错并成功输出 answer；导入 kvpress 后 transformers 中可用 kv-press-text-generation 管道名。

常见问题

注意事项

• 本地源码安装使用 uv（uv sync），不是 pip。
• DecodingPress 为实验性功能，使用前注意稳定性。
• hidden_states_buffer_size 默认 256，部分 press 不需要缓冲隐藏状态时可设为 0。
• 只依据 README 内容，未提及的端口、路径、版本号等请勿自行添加。