大型语言模型（Large Language Models，简称LLMs）正引领着当今科研和开发领域的重大飞跃。研究目标和方法论经历了显著的转变，越来越倾向于以LLMs为核心的方法。然而，这些模型往往伴随着高昂的成本，使得它们在大规模应用中对许多人而言显得遥不可及。因此，降低操作延迟成为了一个严峻的挑战，尤其是在那些需要快速响应的动态应用场景中。
在LLMs中，键值（Key-Value，简称KV）缓存被用于自回归解码。它在推理的预填充阶段存储了多头注意力机制中的键值对。在解码阶段，新的KV对被添加到内存中。KV缓存存储了注意力机制中的中间键和值激活，将计算复杂度从二次方降低到线性。虽然KV缓存提高了效率，但其内存需求随着批量大小、序列长度和模型大小的增加而线性增长。KV缓存的内存需求不断增长，超出了GPU的处理能力，将其转移到CPU引入了多个瓶颈，这增加了延迟并降低了吞吐量。
在将缓存从CPU传输到GPU进行计算的过程中，PCIe接口成为了一个限制因素。慢速的PCIe接口可能导致延迟超过正常水平一个数量级，导致GPU出现大量空闲时间。
以往的研究尝试解决PCIe性能缓慢的问题，但这些方法常常因为数据传输和GPU计算时间不匹配而失败，尤其是在处理大批量和上下文大小的情况下。其他方法依赖于CPU资源，这再次成为一个限制因素。本文讨论了一种新颖的PCIe和GPU优化方法。
南加州大学的研究人员提出了一种高效的CPU-GPU I/O感知型LLM推理方法(https://arxiv.org/abs/2411.17089)，以优化PCIe的利用。这种方法通过部分KV缓存重新计算和异步重叠来解决加载大型KV缓存的系统瓶颈。他们的过程涉及将缓存的较小激活段传输到GPU，而不是传输整个KV缓存。然后GPU从这些较小的激活位重建整个缓存内存。关键在于计算注意力分数，确保信息损失最小。
作者提出了一种完全自动化的方法来确定重新计算和通信分割。这项工作包括三个模块，以最小化GPU延迟：
-
分析器模块：收集系统硬件信息，如PCIe带宽和GPU处理速度。
-
调度器模块：将问题表述为线性规划任务，使用硬件信息和用户配置来确定最优的KV分割点。目标是最大化计算和通信过程之间的重叠。
-
运行时模块：协调两个设备之间的数据传输和管理内存分配。
负责找到最优KV分割的调度器模块有两种工作方式：
逐行调度：通过逐行执行计划减少延迟。在这里，GPU在剩余激活异步加载的同时开始重建KV缓存。
逐列调度：通过在多个批次中重用模型权重来最大化吞吐量，并适应大批量大小的推理。它通过在多个批次中而不是按批次顺序处理每个层，重叠KV缓存和激活的传输与MHA（多头注意力）的计算。进一步使用六进程通信并行策略，运行时模块实现了GPU计算和CPU-GPU通信的并发。
作者使用NVIDIA A100 来做测试，并将GPU通过PCIe 4.0 x16接口连接到CPU，来测试所提出的高效LLM推理框架。实验进行了两个目标，以评估框架的性能：
-
面向延迟的工作负载：所提出的方法优于基线，减少了35.8%的延迟。
-
面向吞吐量的工作负载：该方法相对于基线实现了高达29%的改进。
结论：
CPU-GPU I/O感知型的LLM推理方法在提高LLM推理吞吐量的同时，有效地降低了延迟。该方法通过部分重新计算KV缓存，并将其与数据传输过程相重叠，从而最大限度地减少GPU的空闲时间，增强了整体效率。