Zefan-Cai/KVCache-Factory

解决的问题

KVCache-Factory 解决了长上下文大语言模型（LLM）推理过程中与键值（KV）缓存相关的内存瓶颈。随着上下文长度增加，KV 缓存呈线性增长，消耗大量 GPU 内存并拖慢解码速度。本项目提供了一个统一框架，用于实现、比较和评估各种降低内存占用的策略，同时不牺牲模型性能。

如何工作

该项目作为一个“实验场”，将多种 KV 缓存优化技术整合到一个统一的评估接口下。它支持以下几类方法：

• 压缩与淘汰：移除不重要的 token（例如：StreamingLLM、H2O、SnapKV、PyramidKV）。
• 检索：仅选择与当前查询最相关的 token（例如：Quest、L2Norm）。
• 合并：将多个 KV 对聚合为一个（例如：CAM）。
• 量化：降低缓存值的精度（例如：KIVI、KVQuant、GEAR）。
• 卸载：将 KV 缓存移至 CPU 内存，并逐个流式传输回 GPU，以在不导致 GPU 内存溢出的情况下保持完整精度（HeadInfer）。

包含用于 LongBench、RULER 和「针在 haystack 中」等标准基准的运行器，用于衡量这些方法对准确率和延迟的影响。

适用人群

本工具专为从事 LLM 效率研究的 AI 研究人员和开发者设计，尤其适合希望优化长上下文窗口性能或对比不同 KV 缓存管理策略的人群。

亮点

• 统一接口：支持广泛的基线方法，包括 PyramidKV、SnapKV、H2O 和 StreamingLLM。
• 多样化的优化策略：在一个地方整合了淘汰、检索、合并、量化和无损卸载等多种策略。
• 全面的基准测试：集成支持 LongBench、RULER 和 Needle-in-a-haystack 评估。
• 硬件灵活性：支持大模型（如 Llama-3-70B）的多 GPU 推理，并提供 FlashAttention v2 和 SDPA 两种路径。
• 可视化工具：包含用于可视化注意力模式的工具，帮助理解不同压缩方法如何影响模型关注点。