📢 转载信息
原文链接：https://machinelearningmastery.com/kv-caching-in-llms-a-guide-for-developers/
原文作者：Jason Brownlee
KV 缓存（Key-Value Caching）是现代大型语言模型（LLM）推理的关键优化技术。特别是对于自回归模型，例如用于文本生成的 Transformer 模型，KV 缓存可以显著提高推理速度。本文将深入探讨 KV 缓存的工作原理、其对内存占用的影响以及作为开发人员应如何管理它。
本文将介绍：
- LLM 推理中 KV 缓存的背景。
- KV 缓存的工作原理。
- KV 缓存的内存占用。
- 如何优化 KV 缓存。
LLM 推理中的 KV 缓存
Transformer 模型，如用于文本生成的模型，通常是自回归的。这意味着它们一次生成一个 token，并且生成每个新 token 都依赖于先前生成的所有 token。
假设我们想让模型生成一个包含 100 个 token 的序列。对于序列中的第 100 个 token，模型需要处理所有 100 个 token 作为输入。对于第 101 个 token，模型需要处理所有 101 个 token 作为输入。
如果不使用 KV 缓存，模型必须在每一步都重新计算所有先前 token 的注意力键（Key）和值（Value）向量。这会导致计算成本随着序列长度的增加而呈二次方增长，这在实际应用中是不可接受的。
KV 缓存通过存储先前步骤中计算出的键（K）和值（V）向量，来避免这种重复计算。
KV 缓存的工作原理
Transformer 模型的核心是自注意力机制。在注意力机制中，输入 token 被转换为三个向量：查询（Query）、键（Key）和值（Value）。
对于新生成的每个 token，模型会计算其 Q、K 和 V 向量。然后，查询（Q）向量会与之前所有 token 的键（K）向量进行点积运算，以确定注意力分数。这些分数用于加权值（V）向量，从而产生输出。
KV 缓存的机制如下：
- 首次运行（第一次生成第一个 token）：模型接收输入序列，计算每个 token 的 Q、K、V 向量。
- 缓存 K 和 V：在计算完注意力分数后，模型将所有 token 的 K 和 V 向量存储在一个缓存中。
- 后续运行（生成后续 token）：当生成下一个 token 时，模型只需要计算当前新 token 的 Q 向量。然后，它将这个新的 Q 向量与缓存中存储的 所有先前 token 的 K 向量 进行比较。
- 更新缓存：新 token 的 K 和 V 向量也会被计算并添加到缓存中，供下一步使用。
通过这种方式，KV 缓存避免了在后续生成步骤中重复计算先前 token 的 K 和 V 向量，从而极大地提高了推理效率。
KV 缓存的内存占用
虽然 KV 缓存带来了速度上的巨大提升，但它也显著增加了推理时的内存（VRAM）需求。
KV 缓存的大小主要取决于以下几个因素：
- 批次大小（Batch Size）：同时处理的序列数量。
- 序列长度（Sequence Length）：当前处理的上下文长度。
- 模型维度：模型中隐藏层和注意力头的维度。
对于一个特定的模型层，KV 缓存存储的数据量大约是 $ ext{2} imes ext{BatchSize} imes ext{SeqLen} imes ext{HeadDim} imes ext{NumHeads} imes ext{BytesPerElement}$。
这个公式显示了为什么随着批次大小或序列长度的增加，VRAM 需求会线性增长。在许多部署场景中，KV 缓存占用了大部分可用的 VRAM。
举例来说，对于一个拥有 32 个注意力头的模型，每个头维度为 128，如果我们使用一个批次大小为 1 且序列长度为 1024 的序列，KV 缓存的内存占用是相当可观的。
需要注意的是，KV 缓存不仅存储 K 和 V 向量，还可能包含注意力掩码（attention masks）等辅助数据。
优化 KV 缓存
由于 KV 缓存对性能至关重要，同时也是内存消耗的主要来源，优化它是 LLM 部署的关键。
1. 批处理（Batching）
在标准批处理中，所有输入序列都必须具有相同的长度，并且需要为最长的序列分配足够的 KV 缓存空间。如果序列长度差异很大，会导致内存浪费。
动态批处理（Dynamic Batching）或持续批处理（Continuous Batching）是解决此问题的有效方法。它允许新的请求在旧请求尚未完成时就加入批处理，并且可以更有效地管理 KV 缓存。
2. 量化 KV 缓存
KV 缓存通常以 16 位浮点数（FP16）或 32 位浮点数（FP32）存储。将这些向量量化到更低的精度（如 8 位整数，INT8）可以显著减少内存占用，同时对最终输出质量的影响很小。
例如，将 KV 缓存从 FP16 降到 INT8 可以直接减少 50% 的内存占用。
3. PagedAttention (分页注意力)
PagedAttention 是一种由 vLLM 框架推广的内存管理技术。它借鉴了操作系统中的分页机制，将 KV 缓存分割成固定大小的块（pages），而不是为每个序列预留连续的大块内存。
这种方法可以解决 KV 缓存中的内存碎片化问题，并允许更高效地共享和重用内存空间，从而提高吞吐量。
如上图所示，分页注意力使得不同序列的 KV 块可以存储在内存中的任意位置，极大地提高了内存利用率。
4. 序列和上下文管理
在生成过程中，对于长序列，可以选择性地丢弃旧的或不重要的 token 的 KV 缓存部分，以释放内存或保持上下文长度在可控范围内。
这通常通过滑动窗口注意力或基于重要性的缓存管理策略来实现。
总结
KV 缓存是加速自回归 LLM 推理的基石，通过避免重复计算注意力 K 和 V 向量，将原本二次方的计算复杂度降低到接近线性。
然而，它也带来了显著的 VRAM 消耗。作为开发者，了解如何通过动态批处理、量化、以及像 PagedAttention 这样的高级内存管理技术来优化 KV 缓存的使用，是实现高效、可扩展的 LLM 部署的关键。
🚀 想要体验更好更全面的AI调用？
欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。
评论区