像 GPT-4 这样的 LLMs 在语言理解方面表现出色，但在推理过程中却面临着 GPU 内存使用率高的问题，限制了其在聊天机器人等实时应用中的可扩展性。现有方法通过压缩 KV 缓存来减少内存，但忽略了层间依赖性和预计算内存需求。推理内存的使用主要来自模型参数和 KV 缓存，后者消耗的内存明显更多。例如，一个 70 亿参数的模型使用 14 GB 的参数内存，但 KV 缓存却需要 72 GB。这种巨大的内存需求限制了 LLM 在 GPU 上推理的吞吐量。

来自上海交通大学、小红书和华南理工大学的研究人员开发了 PyramidInfer，通过压缩 KV 缓存来增强 LLM 推理。与忽略层间依赖性和预计算内存需求的现有方法不同，PyramidInfer 逐层仅保留关键的上下文键和值。受最近 tokens 在注意力权重方面一致性的启发，这种方法显著减少了 GPU 内存的使用。实验表明，与现有方法相比，PyramidInfer 将吞吐量提高了 2.2 倍，并将 KV 缓存内存减少了 54% 以上，证明了其在各种任务和模型中的有效性。

高效的策略对于处理不断增长的聊天机器人查询需求至关重要，目标是通过利用 GPU 并行性来最大化吞吐量。一种方法是通过流水线并行和 KV 缓存卸载来增加 GPU 内存，利用多个 GPU 或 RAM。对于有限的 GPU 内存，减少 KV 缓存是另一种选择。像 FlashAttention 2 和 PagedAttention 这样的技术通过优化 CUDA 操作来最小化内存浪费。StreamingLLM、H2O 和 Scissorhands 等方法通过关注最近的上下文或注意力机制来压缩 KV 缓存，但忽略了层差异和预填充阶段的压缩。PyramidInfer 通过在两个阶段都考虑特定于层的压缩来解决这些差距。

推理上下文冗余 (ICR) 和最近注意力一致性 (RAC) 假设的验证启发了 PyramidInfer 的设计。ICR 假设许多上下文键和值在推理过程中是冗余的，并且仅在训练中预测下一个 token 时才需要。使用 40 层 LLaMA 2-13B 模型进行的实验表明，更深的层具有更高的冗余度，允许在不影响输出质量的情况下显著减少 KV 缓存。RAC 确认某些键和值始终受到最近 tokens 的关注，从而能够选择关键上下文 (PVC) 以进行高效推理。PyramidInfer 利用这些见解在预填充和生成阶段有效地压缩 KV 缓存。

PyramidInfer 的性能在各种任务和模型中都得到了评估，证明了在保持生成质量的同时，GPU 内存使用量显著减少，吞吐量也有所提高。评估包括 wikitext-v2 上的语言建模、MMLU 和 BBH 等 LLM 基准测试、使用 GSM8K 进行数学推理、使用 HumanEval 进行编码、使用 MT-Bench 进行对话处理以及使用 LEval 进行长文本摘要。PyramidInfer 在 LLaMA 2、LLaMA 2-Chat、Vicuna 1.5-16k 和 CodeLLaMA 等不同规模的模型上进行了测试。结果表明，PyramidInfer 在比全缓存方法更少的 GPU 内存情况下有效地保持了生成质量，并且明显优于局部策略。

总之，受 ICR 和 RAC 的启发，PyramidInfer 引入了一种在预填充和生成阶段压缩 KV 缓存的有效方法。这种方法在不影响模型性能的情况下显著减少了 GPU 内存使用量，使其成为在资源受限环境中部署大型语言模型的理想选择。尽管 PyramidInfer 非常有效，但它需要额外的计算，限制了小批量情况下的加速。作为第一个在预填充阶段压缩 KV 缓存的方法，PyramidInfer 还没有成为一种无损方法，这表明该领域未来还有改进的潜力。

• 论文地址：https://arxiv.org/abs/2405.12532

⇩ 关注「奇点智源」，探索「人工智能」⇩