arXiv:2411.17089cs.LGcs.DC2024-11ACL被引 16

通过分段重计算与传输,显著降低大模型推理时的I/O延迟。

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

  • 先传部分激活值,让GPU边重算边接收剩余缓存
  • 实测推理延迟降低35.8%,吞吐提升46.2%
  • 自动适配硬件与输入特征,无需人工调参

大语言模型推理计算成本高,依赖键值(KV)缓存存储中间激活以降低生成开销。但缓存内存增长快,常超出GPU容量。将缓存卸载至CPU可缓解内存压力,却受限于CPU-GPU间PCIe带宽。现有方法尝试通过重叠计算与I/O或异构执行来解决,但存在数据移动过多、依赖CPU性能的问题。本文提出KVPR,一种感知I/O的高效推理方法:由CPU优先传输部分激活值,使GPU能立即开始重算部分KV缓存;同时,剩余缓存并行传入。该策略将重算与传输重叠,减少GPU空闲时间。全自动化设计包含分析模块(基于输入特征与硬件信息)、调度模块(优化计算与通信分配)及运行时模块(高效执行计划)。实验表明,相比现有最优方法,KVPR在解码阶段实现最高35.8%的延迟降低和46.2%的吞吐提升。代码已开源。

原文摘要 · Abstract (English)

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.

大模型推理KV缓存I/O优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。