新闻

• [2024-11-28] 更名为 KVCache-Factory！我们项目的目标现在是构建一个统一的框架，用于对多种模型的 KV 缓存进行压缩。

• [2024-06-25] 现已支持使用大语言模型进行多 GPU 推理！快来试试在 LlaMa-3-70B-Instruct 上使用 PyramidKV 吧！

• [2024-06-10] 在 Flash Attention v2 和 Sdpa Attention 中现已支持 PyramidKV、SnapKV、H2O 和 StreamingLLM！如果你的设备（如 V100、3090）不支持 Flash Attention v2，可以将 attn_implementation 设置为 sdpa，从而在 Sdpa Attention 中尝试 PyramidKV！

待办事项：

• 支持 Streaming LLM、H2O 和 SnapKV 的实现
• 支持 Mistral 模型
• 支持 Needle 的实现
• 支持在没有 Flash Attention v2（即 Sdpa Attention）的情况下对 V100 进行 KV 缓存压缩
• 支持 70B LlaMa-3 的多 GPU 推理
• 引入新功能以支持 KV 缓存预算分配（例如百分比支持）
• 支持 Mixtral
• 支持批量推理
• 支持解码阶段的 KV 缓存压缩

性能

可视化：低效注意力

Llama 模型在包含 3 个文档时的注意力图如下所示：

./visualization-tools/vis.ipynb 复现了论文中的可视化结果。我们在 ./visualization 目录下提供了更多可视化工具，支持不同层次的 KV 缓存可视化。 不同层的模型注意力图将存储在 ./attention 目录中。

需求

transformers >= 4.41
flash-attn >= 2.4.0.post1

安装

git clone https://github.com/Zefan-Cai/PyramidKV.git
cd PyramidKV
pip install -r requirements.txt .

推理

我们支持在 LongBench 上运行的推理代码，以复现我们的结果。

请参考 scripts/scripts_longBench/eval.sh，根据你的需求修改参数。

我们的代码库支持 Flash Attention v2、Sdpa Attention 等。我们论文中展示的结果基于 Flash Attention v2。

export CUDA_VISIBLE_DEVICES=$1

method=$2 # 支持 PyramidKV、SnapKV、H2O、StreamingLLM
max_capacity_prompts=64 # 论文中为128、2048
attn_implementation=$3 # 支持 "flash_attention_2"、"sdpa"、"eager"
source_path=$4
model_path=$5
save_dir=${source_path}"results_long_bench" # 结果保存路径

python3 run_longbench.py \
    --method ${method} \
    --model_path ${model_path} \
    --max_capacity_prompts ${max_capacity_prompts} \
    --attn_implementation ${attn_implementation} \
    --save_dir ${save_dir} \
    --use_cache True

• CUDA_VISIBLE_DEVICES：对于大语言模型的多 GPU 推理，只需指定 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7。对于单 GPU 推理，则只需指定 CUDA_VISIBLE_DEVICES=0。
• model_path：你的模型路径。目前支持“Llama-3-8B-Instruct”。
• method：支持 PyramidKV、SnapKV、StreamingLLM、H2O。
• max_capacity_prompts：每层选择的 KV 大小。（例如，论文中为 128 或 2048）。当方法为“PyramidKV”时，由于总的 KV 数量保持不变，每层的具体 KV 长度会相应调整。
• save_dir：保存 LongBench 结果的目录路径。

修改参数后，运行：

sh scripts/scripts_longBench/eval.sh

针尖藏草堆

我们支持在 针尖藏草堆 数据集上运行的推理代码，以复现我们的结果。

请参考 scripts/scripts_needle/eval.sh，根据你的需求修改参数。

我们的代码库支持 Flash Attention v2、Sdpa Attention 等。我们论文中展示的结果基于 Flash Attention v2。

METHOD='pyramidkv'       # ['full', 'pyramidkv', 'snapkv', 'streamingllm', 'h2o']
MAX_CAPACITY_PROMPT=96  # [64, 96, 128, 256, 512, 1024, 2048, ...]
attn_implementation="flash_attention_2" # 支持 "flash_attention_2"、"sdpa"、""。
TAG=test


# 对于 Llama3-8b

(
python -u run_needle_in_haystack.py --s_len 1000 --e_len 8001\
    --model_provider LLaMA3 \
    --model_name /mnt/workspace/zhiyuanhu/yuliang/models/llama3-8b_raw \
    --attn_implementation ${attn_implementation} \
    --step 100 \
    --method $METHOD \
    --max_capacity_prompt $MAX_CAPACITY_PROMPT \
    --model_version LlaMA3_${METHOD}_${MAX_CAPACITY_PROMPT}_${TAG}
) 2>&1  | tee results_needle/logs/LlaMA3_${METHOD}_${MAX_CAPACITY_PROMPT}_${TAG}.log

• LLaMA3 和 Mistral2 的推理都支持单 GPU。
• model_provider：LLaMA3 或 Mistral2
• model_name：你的模型路径。目前支持“Llama-3-8B-Instruct”和“Mistral-7B-Instruct-v0.2”。
• step：上下文长度的增加幅度。
• method：支持 PyramidKV、SnapKV、StreamingLLM、H2O。
• max_capacity_prompt：每层选择的 KV 大小。（例如，论文中为 128 或 2048）。当方法为“PyramidKV”时，由于总的 KV 数量保持不变，每层的具体 KV 长度会相应调整。

要复现我们的结果，请运行：

bash scripts/scripts_needle/eval.sh

推理完成后，运行

python scripts/scripts_needle/visualize.py

以生成图像。你需要将 visualize.py 中的 FOLDER_PATH 修改为你输出的路径（即 eval.sh 中的 --model_version 参数值）。

引用

如果你发现 PyramidKV 对你的研究和应用有帮助，请使用以下 BibTeX 格式引用：

@article{cai2024pyramidkv,
  title={Pyramidkv: 基于金字塔信息漏斗的动态 KV 缓存压缩},
  author={Cai, Zefan 和 Zhang, Yichi 和 Gao, Bofei 和 Liu, Yuliang 和 Liu, Tianyu 和 Lu, Keming 和 Xiong, Wayne 和 Dong, Yue 和 Chang, Baobao 和 Hu, Junjie 和 Xiao Wen},
  journal={arXiv 预印本 arXiv:2406.02069},
  year={2024}
}

@article{fu2024not,
  title={并非所有头都重要：一种集成检索与推理的头部级 KV 缓存压缩方法},
  author={Fu, Yu 和 Cai, Zefan 和 Asi, Abedelkadir 和 Xiong, Wayne 和 Dong, Yue 和 Xiao, Wen},
  journal={arXiv 预印本 arXiv:2410.19258},
  year={2024}
}

致谢

感谢 [SnapKV] SnapKV：LLM 在生成之前就知道你在寻找什么，为我们提供了开源代码，助力本项目的扩展。

KVCache-Factory 快速上手指南

KVCache-Factory 是一个统一的框架，旨在为多种大语言模型提供高效的 KV Cache 压缩方案（如 PyramidKV, SnapKV, H2O, StreamingLLM），支持在单卡或多卡环境下运行，显著降低显存占用并提升长文本推理效率。

环境准备

系统要求

• GPU: 支持 NVIDIA GPU。
  • 若使用 Flash Attention v2（推荐，性能最佳）：需 Ampere 架构及以上（如 A100, RTX 3090/4090）。
  • 若使用 SDPA Attention：兼容旧架构（如 V100, RTX 2080Ti 等）。
• Python: 建议 Python 3.8+
• CUDA: 已正确安装并与 PyTorch 版本匹配。

前置依赖

确保已安装以下基础库（项目会自动安装特定版本）：

• transformers >= 4.41
• flash-attn >= 2.4.0.post1 (可选，但强烈推荐)

国内加速提示：建议使用清华或阿里镜像源加速 pip 安装。

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

安装步骤

1. 克隆项目代码

   git clone https://github.com/Zefan-Cai/PyramidKV.git
   cd PyramidKV
   

2. 安装依赖

   pip install -r requirements.txt .
   

   注：若安装 flash-attn 失败，请确认 CUDA 环境正常，或暂时跳过该依赖使用 SDPA 模式。

基本使用

本项目主要通过脚本在 LongBench 或 Needle in Haystack 基准测试上运行压缩算法。以下以 LongBench 为例展示最简用法。

1. 配置运行参数

你可以直接修改 scripts/scripts_longBench/eval.sh，或通过命令行传入参数。核心参数说明如下：

• method: 压缩算法，支持 PyramidKV, SnapKV, H2O, StreamingLLM。
• max_capacity_prompts: 每层保留的 KV 数量（例如 64, 128, 2048）。
• attn_implementation: 注意力实现方式，支持 flash_attention_2 (推荐), sdpa (兼容旧卡), eager。
• model_path: 本地模型路径（如 Llama-3-8B-Instruct）。

2. 执行推理

假设你使用单张显卡 (ID 0)，运行 PyramidKV 算法，保留 64 个 KV，使用 Flash Attention v2：

export CUDA_VISIBLE_DEVICES=0

method=PyramidKV
max_capacity_prompts=64
attn_implementation=flash_attention_2
source_path=./data/LongBench/ # 请替换为你的数据集路径
model_path=/path/to/Llama-3-8B-Instruct # 请替换为你的模型路径
save_dir=${source_path}"results_long_bench"

python3 run_longbench.py \
    --method ${method} \
    --model_path ${model_path} \
    --max_capacity_prompts ${max_capacity_prompts} \
    --attn_implementation ${attn_implementation} \
    --save_dir ${save_dir} \
    --use_cache True

3. 多卡推理（针对 70B 等大模型）

只需修改 CUDA_VISIBLE_DEVICES 即可自动启用多卡并行：

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
# 其余命令同上

4. 可视化结果 (可选)

对于 Needle in Haystack 任务，推理完成后可生成可视化图表：

# 修改 scripts/scripts_needle/visualize.py 中的 FOLDER_PATH 为你的输出目录
python scripts/scripts_needle/visualize.py