KV 缓存：深度解析大型语言模型推理的内存挑战

大型语言模型（LLM）的推理过程通常需要大量的计算资源，特别是自注意力机制的计算量会随着序列长度的平方增长。为了优化推理效率，KV 缓存应运而生。它通过存储过去token的键值张量来避免重复计算，从而将计算复杂度从平方级降低到线性级。

KV 缓存的权衡：内存换算力

KV 缓存是一种权衡策略，它以牺牲内存为代价来换取更快的计算速度。本文将深入探讨 KV 缓存的规模、面临的挑战以及应对策略。

KV 缓存到底有多大？

对于每个序列中的每个token，我们需要存储两个向量张量（一个键张量和一个值张量），每个张量的维度为 d_head，对应于每个注意力层的每个注意力头的维度。每个张量参数所占用的空间取决于其精度：全精度（FP32）为 4 字节/参数，半精度（BF16、FP16）为 2 字节/参数，8 位数据类型（INT8、FP8）为 1 字节/参数，等等。

假设批次大小为 b，总序列长度（提示 + 完成）为 t，解码器块/注意力层数为 n_layers，每个注意力层的注意力头数为 n_heads，注意力层的隐藏维度为 d_head，精度为 p_a。那么，多头注意力（MHA）模型的 KV 缓存每个token的内存消耗（以字节为单位）为：

$$
KV_cache_per_token = 2 \cdot n_heads \cdot d_head \cdot p_a
$$

注意：在 MHA 模型中，n_heads \cdot d_head = d_model，但为了简化公式，我们不会使用它。

因此，KV 缓存的总大小（以字节为单位）为：

$$
KV_cache_total = b \cdot t \cdot KV_cache_per_token
$$

KV 缓存的挑战：内存爆炸

从公式可以看出，KV 缓存的大小与批次大小和总序列长度成线性关系。由于总序列长度无法事先确定，因此 KV 缓存的内存需求也是未知的，这给内存管理带来了巨大的挑战。

以流行的 MHA 模型为例（表 1），包括 Meta 的 Llama-2 和 OPT，MosaicML 的 MPT 和 BigScience 的 BLOOM：

表 1：流行的多头注意力 (MHA) 模型规格

假设参数存储在半精度（FP16、BF16）中，我们选择一个较小的模型（Llama-2-7B）和一个较大的模型（BLOOM-176B）。对于 Llama-2-7B（BLOOM-176B），KV 缓存的内存消耗约为 0.5 MB/token（4 MB/token）。

以 Llama-2-7B 为例，使用半精度，加载模型权重大约需要 14 GB 的内存，与缓存 28k 个token的键值相同。28k 个token可能对应于 56 个长度为 512 的序列的批次，这并不算极端。

从上面的数字可以看出，KV 缓存的内存消耗会非常大，甚至可能超过加载大型序列模型权重所需的内存量。

应对 KV 缓存挑战：内存优化策略

为了应对 KV 缓存带来的内存压力，我们可以采取多种优化策略：

1. 减少模型权重内存占用

• 权重量化： 使用诸如 AWQ 或 GPTQ 等算法将模型权重量化为更小的数据类型，例如 INT8 或 FP8，从而减少模型权重的内存占用。

2. 减少 KV 缓存内存占用

• 减少批次大小： 降低批次大小可以减少 KV 缓存的内存占用，但也会降低硬件利用率，影响成本效益。
• 减少对总序列长度的依赖：
  • 滑动窗口注意力： 使用诸如 Mistral-7B 等模型，其注意力机制只关注最后几个相邻的token，从而限制了 KV 缓存的大小。
  • 注意力压缩：
    • StreamingLLM 框架： 只保留第一个位置token（“汇聚token”）和最后一个相邻token（局部注意力）的键值，从而构建一个固定长度的滑动窗口。
    • H2O 和 Scissorhands： 通过设置缓存预算并丢弃一些token来压缩 KV 缓存。
    • FastGen： 通过设置最大近似误差来压缩注意力矩阵，从而实现更高的模型精度。
• 减少层数： 较小的模型通常层数较少，如果小型模型能够满足你的用例，可以选择更小的模型。
• 减少注意力头数：
  • 多查询注意力 (MQA)： 所有查询头共享相同的键和值头，从而减少 KV 缓存的大小。
  • 分组查询注意力 (GQA)： 将查询头分成组，同一组内的查询头共享相同的键和值头，提供了一种介于 MHA 和 MQA 之间的折衷方案。
• 减少注意力头隐藏维度： 如果选择相同模型系列的不同大小版本，可能不会对 KV 缓存大小产生影响。
• 使用更少的字节/参数： 对 KV 缓存进行量化可以显著减小其大小，但仅对权重进行量化（如 AWQ 或 GPTQ）无效。需要对权重和激活进行量化（如 LLM.int8() 或 SmoothQuant）才能实现 KV 缓存的量化。

3. 利用多设备内存

• 模型并行： 将模型拆分到多个 GPU 上，利用多个 GPU 的总内存和带宽，从而释放内存压力。
• 内存卸载： 将 KV 缓存卸载到 CPU 内存或磁盘等更廉价的存储设备上，但会带来性能下降。

4. 优化内存管理

• PagedAttention： 将内存分配成固定大小的块，每个块可以包含固定数量的token，并可以跨请求共享，从而减少内存浪费。
• RadixAttention： 通过对 KV 缓存进行索引，实现跨请求的 KV 缓存重用，从而提高性能。

5. 扩展存储容量

• 多 GPU： 使用多个 GPU 可以扩展存储容量，但受到单个主机所能容纳的最大多 GPU 实例的限制。
• 多主机模型并行： 将模型和 KV 缓存拆分到多个主机上，可以实现无限的存储容量和序列长度。

KV 缓存优化：未来方向

随着 LLM 的不断发展，KV 缓存优化将成为一个重要的研究方向。未来，我们将看到更多针对 KV 缓存的优化策略，例如：

• 更有效的注意力机制： 例如，基于稀疏注意力或局部注意力的模型。
• 更智能的缓存压缩算法： 例如，能够根据模型的特定需求动态调整压缩策略的算法。
• 更先进的内存管理技术： 例如，能够根据请求的具体情况动态分配和释放内存的算法。

总结

本文深入探讨了 KV 缓存的内存挑战以及应对策略。通过优化模型权重、KV 缓存、内存管理和存储容量，我们可以有效地解决 KV 缓存带来的内存压力，提高 LLM 推理的效率。

参考文献

[1]: Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron et al., 2023)
[2]: OPT: Open Pre-trained Transformer Language Models (Zhang et al., 2022)
[3]: Release blog posts for: MPT-7B (May 2023) and MPT-30B (June 2023)
[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience, 2023)
[5]: Scaling Laws for Neural Language Models (Kaplan et al., 2020)
[6]: Mistral 7B (Jiang et al., 2023)
[7]: Efficient Streaming Language Models with Attention Sinks (Xiao et al., 2023) + GitHub repository
[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang et al., 2023) + GitHub repository
[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu et al. 2023)
[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge et al., 2023)
[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery et al., 2022)
[14]: The Falcon Series of Open Language Models (Almazrouei et al., 2023)
[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023) + GitHub repository
[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022) + GitHub repository
[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) + GitHub repository
[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022) + GitHub repository
[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng et al., 2023) + GitHub repository
[20] Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023) + GitHub repository
[21] vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon et al. 2023)
[22] Efficiently Programming Large Language Models using SGLang (Zheng et al., 2023) + Blog post
[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang et al., 2018)
[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
[25]: Efficiently Scaling Transformer Inference (Pope et al., 2022)
[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin et al., 2024)