大型语言模型（LLMs）对AI应用产生了深远的影响，尤其是在长文本理解和生成领域。KV缓存（Pope等人，2022年）技术在业界得到了广泛应用，通过缓存先前计算出的KV状态确保了序列生成的效率。然而，这也引入了相当大的内存开销。
作者发现KV缓存并非必要，并提出了新型的KCache技术来缓解LLMs推理过程中的内存瓶颈问题。KCache无需任何训练过程即可直接用于推理。
作者的评估显示，KCache在保持准确性的同时，将流行LLMs的吞吐量提高了40%。
1 Introduction
目前，像GPT-4（OpenAI，2023年）、PaLM 这样的大型语言模型在自然语言处理、总结、代码生成、问题回答等方面占据主导地位。然而，它们昂贵的在线推理成本对基于LLM的应用部署构成了重大障碍。在有限的计算资源下，如何尽可能最大化整个系统的吞吐量，并提高GPU集群的利用率变得越来越重要。
LLM的推理包括两个阶段：prefill 阶段和 decode 阶段。
decode 阶段基于 PREFILL 阶段的结果和上一步骤的 decode 阶段逐一生成 Token ，这一过程受内存带宽的限制。因此，作者需要增加批处理大小以提高系统吞吐量，但增加批处理大小将进一步占用更多的GPU内存。
LLM推理的内存使用主要由个部分组成：模型权重、激活值和KV缓存。例如，对于LLaMA-27B模型，权重在fp16精度下占用大约GB的内存。在处理批大小为和序列长度为时，KV缓存占用大约GB的内存，采用逐层内存共享策略时，激活值仅占用GB的内存。KV Cache，其中表示K缓存和V缓存，表示批大小，表示序列长度，表示嵌入维度，表示层数。随着批大小和序列长度的增加，KV缓存的内存使用将线性增加。
已经提出了一些优化措施来缓解KV缓存瓶颈。量化压缩算法从字节的角度提出压缩KV缓存。上下文窗口压缩算法从s的角度提出压缩KV缓存。自适应计算算法（Schuster等人，2022）通过提前退出解码来减少计算，这是从l的角度。Shazeer（2019年）；Ainslie等人（2023年）通过改进多头注意力（MHA）的结构来加速推理。从K缓存和V缓存的角度来看，尽管在推理过程中简单地卸载到CPU然后再重新加载回GPU可以缓解GPU内存的压力，但当前的宿主到设备（H2D）和设备到宿主（D2H）带宽将成为推理的新瓶颈。
提出仅在推理期间关键 Token 是重要的，因此KV缓存通过删除其中一部分来进行压缩。然而，考虑到多轮问答场景，直接删除KV缓存的一部分而没有回退机制是高风险的操作。一种更灵活的方法是尽可能保留所有KV状态，并动态选择关键信息进行计算。这样，由于保留了所有KV状态，可以保证准确性的上限足够高。
基于这个想法，一个明显的方法是将所有KV状态卸载到CPU内存。另一个关键问题是，如何动态选择哪些KV状态是重要的，并将它们从CPU内存复制回HBM进行注意力计算。
只要这部分信息能最大限度地保留所有语义信息，推理准确性就能尽可能地接近理论上的上限，而部分数据复制也能最大化推理性能。
作者提出了KCache，在推理过程中，作者在HBM中保留K缓存，同时将V缓存存储在CPU内存中。同时，作者直接利用注意力计算得到的softmax结果来筛选关键信息，并从CPU内存中回忆起相应的V缓存以进行后续的注意力计算。通过这个简单的方法，利用Transformer模型的结构特性，作者有效地利用了空闲的CPU内存，增加了HBM的容量。
在本文中，作者基于KCache构建了_InferenceEngine_，该引擎在LLM推理过程中高效地减少了内存占用，实现了40%的吞吐量提升并保持了准确性。作者的工作主要贡献包括：
作者提出了KCache，它无需任何训练过程即可直接用于推理，同时将吞吐量提高了40%，同时保持了准确性。 作者识别了将V缓存卸载到CPU内存时性能和准确性的挑战，提出了KCache来解决这一挑战，并通过模型推理的实验验证了其有效性。 KCache是灵活且可扩展的，可以应用于转换后的预训练模型。
2 Background
在本节中，作者介绍了一些LLM的基本知识，包括自回归推理、预填充和解码。
LLM本质上基于一种仅解码器架构，由个堆叠的块组成。每个块包括两个模块：多头注意力（MHA）（Vaswani et al., 2023）和全连接的前馈网络（FFN）。一个输入张量，其中表示批大小，表示输入的序列长度，表示模型的隐藏维度。
MHA使用个头将输入映射到不同的子空间：，，其中，，，以及，，是可训练的权重，表示每个头的隐藏维度。
。以SwiGLU（Shazeer, 2020）为例，FFN为，其中，，，是激活单元。LLM处理一系列名为提示的单词，并生成一些新单词。
自回归推理意味着当前时刻生成的标记依赖于前一个时刻生成的标记。处理用户提示的过程称为预填充，只需执行一次。在自回归中逐一生成所有输出标记的过程称为解码，需要连续执行。在预填充阶段，以提示作为输入，并使用矩阵-矩阵乘法进行并行计算。在解码阶段，执行与预填充相同的操作，但只接收一个标记作为输入，并使用KV缓存进行向量-矩阵乘法计算。
3 Method
以下是3方法部分的开始。
KCache
在长上下文场景中，用户通常会基于一个长序列提出多轮问题，每一轮问题可能关注长上下文的不同部分。为了最大化每轮结果的准确性，作者避免减少或压缩KV状态，从而确保模型有效性的上限。然而，简单地将KV状态卸载到CPU内存，并在推理时重新加载到GPU上，会显著增加端到端的推理时间。因此，为了平衡模型有效性和推理延迟，作者必须找到一种方法，只重新加载必要的信息回到HBM，这意味着需要一个模块来确定哪些信息是重要的。幸运的是，考虑到注意力机制中键（Key）和值（Value）对的意义，其中键用于与 Query （Query）计算相关性，而值表示与键相关的实际信息，这启发了作者将部分K缓存和V缓存卸载到CPU内存。
图1展示了KCache的方法。作者将K缓存和前层的V缓存保留在HBM中，并将其他V缓存保留在CPU内存中。在计算过程中，注意力计算从调整到。由于K缓存仍然存储在HBM中，的计算不会受到影响。在softmax操作之后，TopN选择最相关的N个结果。作者根据注意力分数动态灵活地将V缓存的相应向量实时移动到HBM，以参与后续计算。
基于所提出的KCache方法，直观上，随着N的增加，模型的推理准确性将接近完整的KV缓存的准确性，但也会增加数据复制开销，导致性能下降。推理性能与推理准确性之间是否存在完美的平衡需要定量分析。在以下各节中，作者将从准确性和性能两个方面提供分析。
Analysis of KCache Performance
在 阶段，需要将 V Cache 的部分异步复制到 CPU 内存中。作者希望每个层的计算时间能够与前一层的数据复制时间重叠。每个 transformer 块需要从设备传输 字节数据到主机，每个 Transformer 块有 次浮点运算（FLOPs）。设：
以 NVIDIA A100 (80GB) GPU 为例，对于 LLaMA2-7B，，这意味着计算将重叠传输。
在 阶段，多头注意力（MHA）模块是一个典型的内存受限任务，正如表1 中所示的算术强度所示。算术强度定义为浮点运算（FLOPs）与 I/O 字节数的比率。这表明在解码过程中，MHA 模块的计算时间强烈依赖于内存访问量。值得注意的是， 阶段 MHA 模块的性能与隐藏大小和序列长度无关，仅受批量大小的影响。
这一观察导致一个预期：所提出的 KCache MHA 模块的计算时间和数据传输时间可以小于传统的 KV 缓存 MHA 实现。设：
以 NVIDIA A100(80GB) GPU（2039GB/s GPU 内存带宽和 32GB/s H2D 带宽）为例，这意味着当 时，KCache 的性能不会降低。
Analysis of KCache Accuracy
在预填充（_prefill_）阶段，将_Value_张量异步卸载到CPU内存，这对推理的准确性和性能没有影响。在解码（_decode_）阶段，有必要减少从主机传输到设备的数据量。基于 ，，其中代表批处理的一个实例，代表一个头（head）。如果的结果足够稀疏，那么对应的值对最终结果的影响将可以忽略不计。在第4节中，将进一步验证KCache的准确性。
Setup
模型与数据集。 所有模型均基于仅解码器的 Transformer 结构，作者在四个开源大型语言模型上评估KCache：LLaMA2-7B、LLaMA2-13B、LLaMA3-8B（AI@Meta，2024）和Mistral-7B 。作者选择了3个基准测试：BBH、GSM8K和TriviaQA。BBH 是一套23个具有挑战性的BIG-Bench任务。GSM8K是一个包含8.5k高质量、语言多样化的中小学数学问题数据集。TriviaQA是一个阅读理解数据集，包含超过650K的问题-答案-证据三元组。
Results
准确性。 为了公平比较，作者使用运行所有任务。
表2展示了关于小样本学习性能的实验结果。
图2显示了不同数据集上的提示长度。
KCache基本上在无损的情况下保持了准确性，甚至在多个数据集和模型上实现了更好的性能。 参数对模型准确性的影响相对较小，尤其是当模型性能足够强时。如果模型性能相对较弱，建议设置较大的。 较大的可以获得更高的准确度。当时，KCache保持了与 Baseline 相同的准确度，甚至更高。作者相信TopN对softmax进行了正则化，并进一步过滤掉噪声信息。 作者在三个数据集上的实验验证了，对于大约2K或更短的上下文长度，将设置为64或128并不会显著影响准确性。
性能。 作者进一步评估了KCache在作者的推理引擎中的端到端性能，并在GPU上进行了实验，该GPU具有64GB内存和1TB的GPU内存带宽以及180TFLOPS。作者评估了LLaMA2-7B。
表3展示了实验结果。总体而言，实验结果进一步验证了3.2中的分析，即当时，KCache显示出性能优势。同时，基于2中的结果，当上下文长度为15K且时，KCache在推理中实现了40%以上的吞吐量提升。
5 Conclusion
在这项工作中，作者提出了KCache，一种针对大型语言模型的高效推理技术。特别是在长上下文推理场景中，KCache实现了40%以上的吞吐量提升。这种方法不需要任何训练，并且适用于各种主流结构，如MHA和GQA。将来，作者将进一步探索基于KCache的策略。
参考
[1].Efficient LLM Inference with KCache.
