Paged Attention 原理#
大模型推理的内存管理挑战#
大模型推理中，处理推理请求的数量受到 GPU 内存容量的限制，也就是说，推理服务系统的吞吐量是 memory-bound，克服这个内存限制需要解决以下内存管理方面的挑战：
KV Cache 内存大。KV Cache 的大小随着请求数量的增加而迅速增长，如图所示，KV Cache 所占内存接近大模型参数的一半。
复杂的解码策略。大模型推理有多种解码策略，每种算法对内存管理的影响各不相同。
输入与输出序列长度未知。在大模型推理中，输入 prompt 与最终输出长度是不确定的。
Paged Attention 原理#
操作系统虚拟内存概念#
<<<<<<< HEAD
Paged Attention 算法#
PagedAttention 算法
PageAttention 允许在非连续的内存空间中存储连续的键和值【与传统 attention 算法不同的是】
将每个序列的 KV 缓存划分为 KV block，每个 block 包含固定数量的 key，value vectors（kv block size）
在注意力计算时，PagedAttention kernel 分别对不同的 KV 块进行识别和 fetch
KV Cache 管理#
=======
Paged Attention 算法#
KV Cache 管理#
!!!!!!!! main
不同解码策略下的表现#
调度与抢占#
LLM 服务面临一个独特的挑战：LLM 的输入提示可能在长度上有很大的变化，而最终的输出长度是先验未知的，取决于输入提示和模型 在此背景下，vLLM 需要回答两个经典问题：( 1 )应该驱逐哪些区块? ( 2 )如果需要，如何恢复被驱逐的块?
通常，驱逐策略使用启发式来预测未来哪个块将被访问得最远，并将该块驱逐出去。由于在我们的情况下，我们知道一个序列的所有块都被访问在一起，因此我们实施了一个全或无的驱逐策略，即要么驱逐一个序列的所有块，要么驱逐一个序列的所有块。此外，一个请求(例如,一个波束搜索请求中的波束候选)中的多个序列被串并作为一个序列组。由于潜在的内存共享，一个序列组内的序列总是被抢占或重新安排在一起。为了回答第二个问题，即如何恢复一个被驱逐的块，我们考虑了两种技术：
Swapping:这是大多数虚拟内存实现所使用的经典技术，它将被删除的页复制到磁盘上的交换空间中。在我们的案例中，我们将被驱逐的块拷贝到 CPU 内存中。如图 4 所示，除了 GPU 块分配器外，vLLM 还包括一个 CPU 块分配器，用于管理交换到 CPU RAM 中的物理块。当 vLLM 为新令牌耗尽空闲物理块时，它选择一组序列驱逐并将它们的 KV 缓存转移到 CPU 上。一旦 vLLM 抢占一个序列并移除它的块，vLLM 就停止接受新的请求，直到所有被抢占的序列都被完成。一旦请求完成，它的块从内存中解放出来，并将抢占序列的块带回来继续该序列的处理。值得注意的是，在本设计中，交换到 CPU RAM 的块数从未超过 GPU RAM 中的总物理块数，因此 CPU RAM 上的交换空间受分配给 KV 缓存的 GPU 内存的限制。
Recomputation：当被抢占的序列被重调度时，我们只需重新计算 KV 缓存。需要注意的是，重新计算的延迟可以显著低于原来的延迟，因为解码时产生的令牌可以与原来的用户提示串接作为新的提示- -它们在所有位置的 KV 缓存可以在一个提示阶段迭代产生。 （理解：比如原始提示只有 5 个 token，但是在它被替换出来的时候已经生成了 3 个 token，那么等它再次计算的时候，我们可以把已经生成的 token 与原始 token 合并起来当作新的 prompt，提示处理阶段是可以并行的，比 decodeing 阶段要快的多）
Comparing Recomputation and Swapping vLLM 同时支持重计算和交换作为其恢复机制。为了理解这两种方法之间的折衷，我们评估了它们的端到端性能，并对它们的开销进行了微基准测试，如图 19 所示。我们的结果表明，在块大小较小的情况下，交换会带来过多的开销。这是因为较小的分块大小往往导致 CPU 和 GPU 之间存在大量的小数据传输，限制了 PCIe 的有效带宽。相比之下，重计算的开销在不同的块大小之间保持不变，因为重计算没有使用 KV 块。因此，当块尺寸较小时，重计算更有效；当块尺寸较大时，交换更有效，尽管重计算开销较大
分布式场景内存管理#
许多 LLM 的参数规模超过了单个 GPU 的处理能力,需要将它们划分到分布式 GPU 上,并以模型并行的方式执行,这就需要一个能够处理分布式内存的内存管理器.
我们观察到，即使模型并行执行，每个模型分片仍然处理相同的输入令牌集合，因此需要 KV Cache 为相同的位置。因此，vLLM 在集中式调度器中具有单个 KV 缓存管理器，如图 4 所示。不同的 GPU 工作者共享管理器，以及从逻辑块到物理块的映射。这种通用映射允许 GPU 工作者使用调度器为每个输入请求提供的物理块来执行模型。虽然每个 GPU 工作者具有相同的物理块 ID，但是一个工作者只存储其对应注意力头的一部分 KV 缓存。
Attention 算子适配#
由于 PagedAttention 引入了现有系统无法有效支持的访存模式，我们开发了多个 GPU 内核对其进行优化：
融合整形和分块写入。在每个 Transformer 层中，新的 KV cache 被分割成块，重新配置为块读取优化的内存布局，然后保存在块表指定的位置。为了最小化内核启动开销，我们将它们融合到单个内核中。
融合块阅读和注意力。将 FasterTransformer 中的注意力核改编为根据块表读取 KV 缓存，并进行动态的注意力操作。为了保证合并后的内存访问，我们为每个块分配一个 GPU 扭曲来读取。此外，我们还增加了对请求批次内可变序列长度的支持。
融合块拷贝。块拷贝操作，由写时拷贝机制发布，可以在不连续的块上操作。如果使用 cudaMemcpyAsync API，这将导致大量的小数据移动的调用。为了减少开销，我们实现了一个内核，将不同块的复制操作批处理到一个内核启动中。