KVPR: Efficient LLM Inference with I/O-Aware KV Cache Partial Recomputation

📄 arXiv: 2411.17089v2 📥 PDF

作者: Chaoyi Jiang, Lei Gao, Hossein Entezari Zarch, Murali Annavaram

分类: cs.LG, cs.DC, cs.PF

发布日期: 2024-11-26 (更新: 2025-06-04)

备注: ACL Findings 2025

🔗 代码/项目: GITHUB

💡 一句话要点

KVPR：一种I/O感知的KV缓存部分重计算方法，用于高效LLM推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 LLM推理 KV缓存 I/O优化 部分重计算 GPU加速 PCIe带宽

📋 核心要点

  1. 现有LLM推理方法在KV缓存过大时，依赖CPU内存导致PCIe带宽成为瓶颈，且难以完全隐藏通信延迟。
  2. KVPR的核心思想是利用GPU的部分KV缓存重计算与CPU的KV缓存传输并行执行，从而减少GPU空闲时间。
  3. 实验表明，KVPR相比现有方法，在解码阶段实现了显著的延迟降低和吞吐量提升，最高分别达到35.8%和46.2%。

📝 摘要（中文）

大型语言模型（LLM）的推理计算量巨大。为了降低自回归解码的成本，Key-Value（KV）缓存用于存储中间激活值，从而显著降低了token生成的计算开销。然而，KV缓存所需的内存迅速增长，通常超过GPU内存的容量。一种经济高效的替代方案是将KV缓存卸载到CPU内存，这减轻了GPU内存压力，但瓶颈转移到CPU和GPU之间PCIe连接的有限带宽。现有方法试图通过将GPU计算与I/O重叠或采用CPU-GPU异构执行来解决这些问题，但它们受到过度数据移动和依赖CPU能力的阻碍。随着KV缓存大小的增长和/或GPU计算能力的提高，完全重叠PCIe通信延迟变得具有挑战性。本文介绍了一种高效的I/O感知LLM推理方法KVPR，其中CPU首先传输一部分激活值，GPU可以从中开始重新计算KV缓存值。当GPU重新计算部分KV缓存时，KV缓存的剩余部分同时从CPU传输。这种方法将GPU重新计算与KV缓存传输重叠，以最大限度地减少GPU空闲时间并最大限度地提高推理性能。KVPR通过集成一个利用输入特征和系统硬件信息的分析器模块、一个调度器模块来优化计算和通信工作负载的分配以及一个运行时模块来高效执行导出的执行计划来实现完全自动化。实验结果表明，与最先进的方法相比，KVPR在解码期间实现了高达35.8%的延迟降低和46.2%的吞吐量提高。

🔬 方法详解

问题定义：论文旨在解决大型语言模型推理过程中，当KV缓存超出GPU内存容量时，将KV缓存卸载到CPU内存所导致的PCIe带宽瓶颈问题。现有方法如GPU计算与I/O重叠或CPU-GPU异构执行，由于数据移动开销大和依赖CPU计算能力，无法有效隐藏通信延迟，导致GPU空闲时间过长。

核心思路：KVPR的核心思路是利用GPU的计算能力，对部分KV缓存进行重计算，并将其与CPU到GPU的KV缓存传输并行执行。通过这种方式，GPU可以在等待剩余KV缓存传输的同时进行计算，从而减少GPU的空闲时间，提高整体推理效率。

技术框架：KVPR包含三个主要模块：分析器模块、调度器模块和运行时模块。分析器模块负责分析输入特征和系统硬件信息，为后续的调度提供依据。调度器模块根据分析结果，优化计算和通信工作负载的分配，生成执行计划。运行时模块则负责高效地执行调度器生成的执行计划，实现KV缓存的部分重计算和并行传输。

关键创新：KVPR的关键创新在于I/O感知的KV缓存部分重计算策略。与现有方法不同，KVPR不是简单地将整个KV缓存卸载到CPU，而是选择性地传输一部分KV缓存，并利用GPU进行重计算，从而实现了计算和通信的有效重叠。

关键设计：KVPR的关键设计包括：(1) 分析器模块如何准确预测计算和通信开销；(2) 调度器模块如何优化计算和通信任务的分配，以最大化GPU利用率；(3) 运行时模块如何高效地执行部分KV缓存的重计算和并行传输。论文中可能包含关于这些模块的具体算法和参数设置的描述，但具体细节未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

KVPR通过I/O感知的KV缓存部分重计算策略，显著提升了LLM推理的性能。实验结果表明，与现有最先进的方法相比，KVPR在解码阶段实现了高达35.8%的延迟降低和46.2%的吞吐量提高。这些结果验证了KVPR在提高LLM推理效率方面的有效性。

🎯 应用场景

KVPR可应用于各种需要高效LLM推理的场景，例如在线对话系统、文本生成服务和机器翻译等。通过降低推理延迟和提高吞吐量，KVPR能够提升用户体验，并降低部署LLM的成本。该研究对于推动LLM在资源受限环境下的应用具有重要意义。

📄 摘要（原文）

Inference for Large Language Models (LLMs) is computationally demanding. To reduce the cost of auto-regressive decoding, Key-Value (KV) cache is used to store intermediate activations, which significantly lowers the computational overhead for token generation. However, the memory required for the KV cache grows rapidly, often exceeding the capacity of GPU memory. A cost-effective alternative is to offload KV cache to CPU memory, which alleviates GPU memory pressure, but shifts the bottleneck to the limited bandwidth of the PCIe connection between the CPU and GPU. Existing methods attempt to address these issues by overlapping GPU computation with I/O or employing CPU-GPU heterogeneous execution, but they are hindered by excessive data movement and dependence on CPU capabilities. Fully overlapping PCIe communication latency gets challenging as the size of the KV cache grows and/or the GPU compute capabilities increase. In this paper, we introduce KVPR, an efficient I/O-aware LLM inference method where the CPU first transfers a partial set of activations, from which the GPU can start recomputing the KV cache values. While the GPU recomputes the partial KV cache, the remaining portion of the KV cache is transferred concurrently from the CPU. This approach overlaps GPU recomputation with KV cache transfer to minimize idle GPU time and maximize inference performance. KVPR is fully automated by integrating a profiler module that utilizes input characteristics and system hardware information, a scheduler module to optimize the distribution of computation and communication workloads, and a runtime module to efficiently execute the derived execution plan. Experimental results show that KVPR achieves up to 35.8% lower latency and 46.2% higher throughput during decoding compared to state-of-the-art approaches. The code is available at https://github.com/chaoyij/KVPR.