LLM推理优化笔记2: vLLM原理PagedAttention

PagedAttention 和vLLM

LLM进行推理时，因为每一个请求的KV cache内存很大（比如PagedAttention论文图1示意13B LLM推理时KV cache存储占到了30%以上的显存）、并且会动态增加或者减小；所以如果KV cache管理不当，会因为内存碎片化和冗余KV缓存造成大量的内存浪费（如论文图2示意）。

受操作系统用带分页机制的虚拟内存(virtual memory with paging)来解决内存碎片化和内存共享的启发，vLLM团队提出PagedAttention算法来更有效管理KV cache，并基于PagedAttention构建了vLLM推理系统。

给定一个prompt，LLM推理服务的生成过程分为两步：

1. prefill(PagedAttention论文中称为The prompt phase)：将prompt 的tokens$(x_1,\dots ,x_n)$作为输入，计算第一个新token的概率$P(x_{n+1}|x_1,\dots ,x_n)$。在这个过程中会计算出key向量 $k_1,\dots ,k_n$和value向量$v_1,\dots ,v_n$，因为全部prompt tokens $x_1,\dots ,x_n$是已知的，所以这个阶段可以通过矩阵乘法(matrix-matrix multiplication)操作进行并行计算。 故这个阶段是compute-bound的。

2. decoding(PagedAttention论文中称为The autoregressive generation phase)：按顺序一个接一个地生成新的token，直到生成序列达到最大长度(由用户指定或者受限于LLM的上下文长度)或者遇到终止符(end-of-sequence()token)。在第t次迭代生成时，模型以token $x_{n+t}$作为输入计算概率$P(x_{n+t+1}|x_1,\dots ,x_{n+t})$，过程中第1个位置到$n+t-1$个位置的key向量和value向量是在之前的计算迭代中缓存的，只有key向量$k_{n+t}$和$v_{n+t}$是在这次迭代中计算的。这个阶段的计算因为每一次计算都依赖前一次迭代的结果，无法并行计算只能使用matrix-vector multiplication，是memory-bound的。

PagedAttention

一般而言，LLM推理过程中的内存管理有如下挑战：

• Large KV cache，KV cache的大小随着请求次数迅速增加。
• Complex decoding algorithm，存在多种不同的解码算法使得内存管理更复杂，尤其在需要考虑KV cache共享时。
• Scheduding for unknown input & output lengths，用户请求的输入和输出长度都是不确定的，当输出长度增加时，KV cache也会增加极端情况下甚至会用光可用内存。

为了解决上述内存挑战，受操作系统里的分页机制(paging)启发，PagedAttention允许在非连续内存空间存储连续的key和value向量。

PagedAttention将每一个序列的KV cache划分为KV blocks，每一个block包含固定数目的token数（称为KV block size(B)），在vLLM中默认的block size为16。设key block 记为$K_j=(k_{(j-1)B+1},\dots ,k_{jB})$，value block记为$V_j=(v_{(j-1)B+1},\dots ,v_{jB})$，则attention计算可通过如下的block-wise计算来完成，式中的$A_{ij}=(a_{i,(j-1)B+1},\dots ,k_{i,jB})$是在第j个block上attention分数行向量。
$$A_{ij}=\frac{\exp \left(q_i^{\top }{K}_j/\sqrt{d}\right)}{{\sum }_{t=1}^{\lceil i/B\rceil }\exp \left(q_i^{\top }{K}_{t}1/\sqrt{d}\right)},o_i=\sum _{j=1}^{\lceil i/B\rceil }{V}_j{A}_{ij}^{\top },$$
在attention计算过程中，PagedAttention算法独立地识别和获取不同的KV blocks，示例如上图。

vLLM

基于PagedAttention的vLLM架构如下图，Scheduler协调分布式GPU worker的执行，KV cache manager基于PagedAttention来管理KV cache，其通过scheduler的指令来管理GPU worker上的物理KV cache内存。

KV Cache Manager

按照PagedAttention的思想，将KV cache划分为固定尺寸的KV blocks。将KV blocks分为logical KV blocks和pysical KV blocks。

将一个请求的KV cache表示为一系列的logical KV blocks，它是连续的，按token顺序从左到右填充KV block，只有最后一个KV block的空闲位置是保留给将来生成的token使用的。

block magager将GPU worker上的GPU DRAM或者CPU RAM划分为pysical KV blocks。

维护block tables来记录每个请求的logical KV blocks和physical KV blocks之间的映射，block table的每一行记录logical block对应的physical blocks以及填充位置个数。

在vLLM 0.1.0版代码中(早期版本的vLLM更适合了解主要思路和实现）：

• 类LogicalTokenBlock 定义了logical block，有block_number/block_size/num_tokens 属性，以及存储了token_ids。

• 类PhysicalTokenBlock 存储KV cache的物理block状态，有device/block_number/block_size/ref_count 属性，如果ref_count为0，这块内存就是空闲的。

• 类BlockAllocator 维护和分配空闲的物理blocks。

• 类BlockSpaceManager 管理逻辑blocks和物理blocks之间的映射，定义了gpu BlockAllocator和cpu BlockAllocator。

Decoding with PagedAttention and vLLM

以上图为例来说明vLLM对单个输入序列解码时是如何执行PagedAttention和管理内存的：

1. 进行prefill步骤，vLLM只会保留在prompt计算过程所需要的KV blocks。在示例中KV block size为4，prompt有7个token，vLLM为其分配了2个logical KV blocks(0和1)，对应的physical block为7和1。在logical block 1还剩下一个槽位，block table的#filled记录的是当前block已存储token的个数。
2. 进行decoding步骤的第一次迭代生成新token，使用physical block 7和1中存储的KV cache进行PagedAttention计算，并将新生成的KV cache存储到logical block 1，相应地更新#filled的值。
3. 进行decoding步骤的第二次迭代生成新token，因为上一个logical block已经满了，将新生成的KV cache存储到logical block 2，并分配新的physical block 3，在block table中记录对应的映射关系。

总之，vLLM随着LLM生成过程动态地分配physical block，只有当前面的block已满才会分配新的physical block，所以可以更有效地利用内存。当一个请求的生成已经结束时，其对应的KV blocks可以被释放掉供其他请求使用。

下图是当有两个请求时，vLLM如何管理内存的示意图，两个请求的physical GPU不需要是连续的。

Parallel sampling

有时候对于单个prompt输入，LLM需要生成多个采样输出，这样用户能够选择一个最喜欢的输出。在这种场景下，一个请求对应的多个输出共享相同的prompt，vLLM通过其PageAttention和分页内存管理，可以很容易实现输入prompt的KV cache共享从而节省内存。

上图是一个有两个输出的并行解码(parallel decoding)例子：

1. 在prefill阶段，只存储一份prompt对应的KV cache，两个请求对应的logical KV blocks指向相同的两个physical KV block。注意因为一个physical block可以对应到多个logical blocks，所以每个physical block都有一个reference count属性。所以在这个例子中physical block 7和1的reference count都为2。
2. 在decoding阶段，两个输出采样了不同的token后需要不同的KV cache存储，vLLM实现了一个physical block粒度的copy-on-write 机制，比如当sample A1需要对其logical block 1进行写入新的token KV cache时，vLLM发现其对应的physical block 1的reference count大于1，它会为sample A1的logical block 1分配一个新的physical block（此例为physical block 3)，并将physical block 1的信息拷贝到physical block 3，同时将physical block 1的reference count减1。而当sample A2开始写入数据到physical block 1时，因为其reference count已经为1了，A2就可以直接将其新生成的KV cache直接写入到physical block 1了。

通过这个例子，我们知道vLLM通过copy-on-write机制，可以在多个输出样本之间共享大部分prompt的KV cache，这对于长输入prompts(long input prompts)生成多个输出时可以节省很多的内存占用。

Beam search

Beam search是在机器翻译等常用的解码方法，它依赖于beam width参数k，用来定义在每一步时top k候选词。beam search不同的输出不仅会共享初始prompts blocks，也会在生成过程中共享其他的block，这个共享模式在生成过程中是动态变化的。

在上图的一个beam search示例中，VLLM管理k=4的beam search的KV blocks，在图中的虚线之前，每一个候选序列有4个logical blocks，所有序列都共享第一个block 0，Beam candidate 3在第二个block开始与其他候选序列不一样，Beam candidate 0-2共享前3个block在第四个block开始不一致。在接下来的迭代生成中，top-4候选项都由Beam candidate 1和2生成，因为原来的Beam candidate 0和3不再属于top候选序列，它们的logical blocks被释放掉，对应的physical blocks的reference count也减少。vLLM因此释放掉所有reference count为0的physical blocks（block 2,4,5,8)，并分配新的physical blocks（block 9-12）来存储新的KV cache。

Shared prefix

一般在使用LLM时，用户会使用包含了任务指令、输入输出样例等的system prompt，system prompt会作为最终用户输入给LLM的prompt的一部分，也就意味着许多请求都会共享一个前缀。在vLLM中会保留预先定义的共享前缀对应的physical blocks，这样当用户请求的prompt包含共享前缀时只需要将其logical blocks映射到已经缓存好的physical blocks就可以了（最后一个block需要copy-on-write机制），可以省略掉共享前缀部分的计算。

Scheduling and Preemption

对来到系统的请求，vLLM执行先进先出(first-come-first-serve, FCFS)的调度策略，当内存不够时，保证先处理最先到达系统的请求，先将最近到达系统的请求prempted。

没有想到用什么词来翻译preempt，在论文的意思大概是因为内存不够将请求延迟处理，如果一个请求已经占用了内存，但需要延迟处理时，将其内存先释放掉

在前面提到内存管理挑战时说到LLM服务请求的长度是变化，当请求数量和请求输出长度不断增加时，vLLM会因为存储新生成的KV cache而耗光GPU phycical block（即GPU内存）。此时vLLM面临的两个经典问题是：1. 需要将那些已经占用的blocks驱逐掉？2.如果需要这些被驱逐掉的blocks时如何恢复呢？

对于问题1，vLLM执行的是all-or-nothing驱逐策略，也就是对一个待驱逐请求的所有blocks都进行驱逐，因为每次这些blocks是同时被访问的。同时因为vLLM对于一个请求会生成多个输出序列(比如beam search请求的beam candidates)一起作为一个序列组(sequence group)来调度，因为这些序列会共享内存，所以一个序列组的序列总是被同时preempted或者重新调度。

对于问题，vLLM有两种技术来实现：

• Swapping，将被preempted的序列对应的blocks的KV cache放到CPU。如果有序列被preempted，vLLM就会停止接收新的请求直到所有被preempted的序列也结束后再接收新的请求。当一个正在处理的序列完成后，其blocks被释放，在CPU中存储的序列的block被转换到GPU继续处理。
• Recomputation：重新计算被preempted的序列。因为这些序列已经生成了部分了token，可以将它们直接添加到请求输入prompt中，所以重新计算的速度会比直接处理原始prompt要快。

论文通过实验表明对于较小的block size, Recomputing速度更快，对于较大的block size，Swapping更快，而中等大小的block size(16-64)两者速度差不多。

Distributed Execution

vLLM通过支持Metatron-LM风格的模型并行来支持分布式设置。

如前面vLLM架构图所示意，它只有一个中央调度器，不同GPU worker共享KV cache管理器。