1.论文名称: Efficient Memory Management for Large Language Model Serving with PagedAttention
2.论文链接: https://arxiv.org/abs/2309.06180
3.github项目: https://github.com/vllm-project/vllm.git
4.团队: UC Berkeley
目录:
1. LLM推理面临的内存挑战
1.1 大的KV缓存
1.2 复杂的解码算法(decoding)
1.3 未知输入和输出长度的调度(scheduling)
1.4 现有系统中的内存管理问题
2. 解决措施
2.1 分页注意力【PagedAttention】
2.1.1 背景
2.1.2 实现细节
2.1.3 注意力计算中的具体操作
2.2 KV缓存管理器【KV Cache Manager】
2.2.1 虚拟内存的原理
2.2.2 vLLM的应用
2.2.3 GPU worker上的块引擎
2.2.4 KV块管理器
2.2.5 逻辑与物理的分离
2.3 基于分页注意力和vLLM的解码策略【Decoding with PagedAttention and vLLM】
2.3.1 虚拟内存与预留
2.3.2 第一步自回归解码
2.3.3 第二步解码
2.3.4 解码迭代
2.3.5 动态内存管理
2.3.6 物理与逻辑块的映射
2.4 调度与抢占【Scheduling and Preemption】
2.4.1 背景
2.4.2 如何恢复被驱逐的块，并介绍了两种核心技术
（1）交换 Swapping
（2）重计算 Recomputation
2.5 分布式执行【Distributed Execution】
3. 实施
3.1 内核级优化
3.1.1 融合重塑和块写入（Fused reshape and block write）
3.1.2 融合块读取和注意力（Fusing block read and attention）
3.1.3 融合块复制（Fused block copy）
3.2 支持各种解码算法
3.2.1 Fork 方法
3.2.2 Append 方法
3.2.3 Free 方法
1. LLM推理面临的内存挑战
当我们进行微批处理(mini-batch)时，虽然能减少计算浪费并以更灵活的方式批处理请求，但由于GPU内存容量的限制（特别是存储 KV 缓存的空间），仍然限制了可以一起批处理的请求数量，这意味着服务系统的吞吐量受到内存的限制。具体的内存管理挑战有如下三个方面:
1.1 大的KV缓存
随着请求数量的不断增加，KV缓存的大小也在迅速扩大，这在处理大规模数据时尤为明显。以一个含有13B参数的OPT模型为例，每一个token的KV缓存就需要占用800 KB的空间。而OPT模型能生成的token序列最多可达2048个，因此在处理一个请求时，KV缓存所需的内存空间可能高达1.6 GB。这种情况在当前GPU的资源稀缺环境下尤为突出，因为即便是主流GPU的内存容量也只有几十个GB，如果将所有可用内存都分配给KV缓存，那么也仅能处理几十个请求。而且，如果内存管理不够高效，还会进一步降低批处理的大小，导致资源利用率进一步降低。与此同时，GPU的计算速度的增长速度是超过内存容量的，这让我们相信，随着时间的推进，内存的瓶颈问题将变得越来越明显，可能会严重影响数据处理和模型训练的效率。
1.2 复杂的解码算法(decoding)
LLM服务提供了多种解码算法供用户选择，每种算法对内存管理的复杂性都存在不同的影响。以多个随机样本请求为例，当用户从单个输入提示中请求多个随机样本时，可以通过共享prompt部分的KV缓存来最小化内存的使用。然而，在自回归生成阶段，由于不同样本结果及其相关的上下文和位置依赖关系，保持KV缓存的独立是必要的，以避免数据混淆或错误的生成。
解码算法的不同选择会直接影响到KV缓存共享的程度。例如，在使用较为复杂的解码算法如beam search时，不同的request beams可以共享更大部分的KV缓存，同时，这种共享模式还会随着解码过程的进行而发生变化。这样的设计允许在保证生成质量的同时，尽可能地降低内存使用，从而在一定程度上缓解了内存管理的复杂性，优化了系统的性能表现。通过合理的内存管理和解码算法选择，可以在满足用户需求的同时，保持系统运行的高效和稳定。
1.3 未知输入和输出长度的调度(scheduling)
LLM服务的请求在输入和输出长度上的可变性提出了对内存管理系统的特殊要求，即需要能够适应不同长度的提示。随着请求输出长度的增加，所需的KV缓存的内存也会相应增加，这可能会消耗掉针对新请求或现有提示的进行中生成的可用内存。因此，系统需要进行相应的调度决策以应对这种情况。例如，可以从GPU内存中删除或交换某些请求的KV缓存，以释放内存空间来处理新的请求或继续现有的生成任务。这样的调度决策有助于确保内存资源的有效利用，同时保持LLM服务的高效运行，满足不同长度请求的处理需求。
1.4 现有系统中的内存管理问题
在当前的深度学习框架中，通常要求将张量（tensors）存储在连续的内存区域中。依此，过去的LLM服务系统也采取了相似的做法，将一个请求的KV缓存作为一个连续的张量进行存储。由于从LLM得到的输出长度具有不确定性，这些系统通常会根据一个请求可能的最大序列长度为其静态分配一块内存，而不考虑请求的实际输入或最终的输出长度。这种预分配的策略导致了三个主要的内存浪费来源：
(1) 未来tokens的预留slots：系统预先为可能生成的tokens保留了内存，尽管最终可能不需要所有这些内存。
(2) 内部碎片化：由于为潜在的最大序列长度预留了过多的空间，导致内部的内存浪费。
(3) 外部碎片化：由于采用了如buddy allocator这样的内存分配器，会产生外部碎片化的问题。这些外部碎片化的内存将永远不会被用于生成的tokens，这点在服务请求之前就已经知晓。
尽管这些预留的内存最终可能会被使用，但是为整个请求的持续时间预留这些空间，特别是当预留的空间很大时，占据了本可以用于处理其他请求的宝贵空间。
虽然有人提出使用压缩（compaction）作为解决碎片化问题的一个可能方案，但在一个对性能敏感的LLM服务系统中实施压缩是不切实际的，因为KV缓存的大小非常大。即便采用了压缩，每个请求预分配的块空间也会阻止现有内存管理系统中针对解码算法的内存共享，从而可能对系统的效率和性能产生负面影响。在这种情况下，可能需要探讨更加灵活和高效的内存管理和分配策略，以便在保证LLM服务性能的同时，最大限度地减少内存的浪费。
这个图表(Figure-2)显示了不同的LLM服务系统在实验中的内存浪费的平均百分比。
(1) 颜色解释：
- 绿色：Token状态。这是实际用于存储token状态的内存部分。
- 橙色：预留。这是预先保留但未使用的内存部分。
- 红色：内部碎片。这是由于内存分配造成的未使用的内存部分。
- 灰色：外部碎片与其他。这是由于系统其他部分引起的未使用的内存。
(2) 系统对比：
- Orca (Max)：在这个配置中，Token状态占用了20.4%的内存，预留了57.3%，内部碎片是13.3%，而外部碎片与其他占用了8.9%。
- Orca (Pow2)：在这个配置中，Token状态占用了26.8%的内存，预留了17.9%，内部碎片是13.6%，而外部碎片与其他占用了41.6%。
- Orca (Oracle)：在这个配置中，Token状态占用了38.2%的内存，预留了25.2%，内部碎片是36.6%。
- vLLM：vLLM系统在Token状态上使用了96.3%的内存，几乎没有其他浪费。这显示了vLLM在内存管理上的高效性。
(3) 结论：
从图中可以看出，vLLM系统在内存使用上非常高效，几乎所有的内存都用于Token状态，而浪费非常少。相比之下，Orca系统（无论是Max、Pow2还是Oracle配置）都有不同程度的内存浪费，尤其是预留和碎片化。
2. 解决措施
vLLM这个系统采用一个集中式的调度器来协调分布式的GPU工作节点。其中的KV缓存管理器能够以“分页”的方式有效地管理KV缓存，这正是得益于PagedAttention算法。具体来说，KV缓存管理器通过中央调度器发送的指令来管理GPU工作节点上的物理KV缓存内存。
这是vLLM系统的概览图。现在，让我们一步步理解这个图：
(1) Scheduler (调度器)
一个中央组件，负责协调系统中的其他组件。从图中可以看到，它直接与各个Worker进行交互。
(2) KV Cache Manager (键值缓存管理器)
此管理器控制键值缓存，它的设计允许以分页的方式管理内存。它有两个表现形式的“块表”（Block tables）。
(3) CPU Block Allocator & GPU Block Allocator (CPU块分配器 & GPU块分配器)
这两个组件分别负责为CPU和GPU分配内存块。这种分块的分配方式可能是为了优化内存使用和提高整体系统的效率。
(4) Workers (工作节点)
每个工作节点都由以下部分组成：
- Cache Engine (缓存引擎)：可能用于快速存取数据，使模型运行更加高效。
- Model Shard (模型分片)：这表明模型被分成了多个分片，每个工作节点只处理其中的一个分片。这是一种常见的方法来分布式地处理大型模型，因为它允许多个GPU同时工作，每个GPU只处理模型的一部分。
从这个系统概览中，我们可以得到以下几点关键认识：
- vLLM系统利用了分布式处理来处理大型语言模型。 - 它有一个集中的调度器来协调各个组件。 - 该系统使用了分页的方式来管理内存，为了优化内存使用。 - 通过将模型分片到多个工作节点上，系统可以并行处理任务。
2.1 分页注意力【PagedAttention】
2.1.1 背景
传统的注意力算法往往要求在连续的内存空间中存储键和值，但PagedAttention允许在非连续的内存空间中存储连续的键和值。
2.1.2 实现细节
（1）PagedAttention将每个序列的KV缓存分割成多个KV块。
（2）每个块包含固定数量tokens的键和值向量，这个固定的数量被称为KV块大小（
