PyramidKV

介绍PyramidKV项目

PyramidKV是一个创新的项目，专注于动态KV缓存压缩技术，旨在提高大规模语言模型（LLM）的推理效率和性能。这个项目结合了金字塔信息漏斗的概念，使得在处理庞大的数据集时能够更加有效率地执行运算。

项目待办事项

PyramidKV项目正在进行持续开发，目前已完成以下任务：

支持Streaming LLM、H2O和SnapKV的实现。
兼容Mistral模型。
支持Needle实现。
在不依赖Flash Attention v2的情况下支持KV缓存压缩，使其在V100等设备上运行。
支持70B LlaMa-3的多GPU推理。

未来的开发计划包括：

添加新功能以支持KV缓存的预算分配（如支持百分比）。
支持Mixtral。
支持批量推理。
支持在解码阶段的KV缓存压缩。

项目性能

PyramidKV项目通过一系列图表展示了项目在不同情况下的推理性能和效率提升。图中展示的结果证明了该项目在处理大型语言模型任务时的优势。

注意力的可视化

该项目还提供了可视化工具来分析模型的注意力图。例如，对于Llama模型，可以生成多个层级的注意力图，以帮助理解模型如何处理多个文档。

环境要求

为了运行PyramidKV，需要安装以下依赖：

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

安装步骤

要开始使用PyramidKV，可以通过以下步骤进行安装：

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

推理

项目支持在LongBench上进行推理以重现项目中的结果。用户可以通过调整eval.sh脚本中的参数来符合自己的需求。支持的注意力实现包括Flash Attention v2和Sdpa Attention，项目结果基于Flash Attention v2。

对于大规模LLM的多GPU推理，用户需设置合适的环境变量。例如，CUDA_VISIBLE_DEVICES用于指定使用的GPU设备。本项目目前支持的模型包括Llama-3-8B-Instruct。

针对Needle in Haystack的推理支持

PyramidKV也支持在Needle in Haystack上进行推理，用户可以根据需要修改相应的参数并运行相应的脚本来获得实验结果。

引用

如果PyramidKV对您的研究或应用有所帮助，请使用以下BibTeX格式进行引用：

@article{cai2024pyramidkv,
  title={PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling},
  author={Cai, Zefan and Zhang, Yichi and Gao, Bofei and Liu, Yuliang and Liu, Tianyu and Lu, Keming and Xiong, Wayne and Dong, Yue and Chang, Baobao and Hu, Junjie and others},
  journal={CoRR},
  year={2024}
}