一.背景

大量LLM应用(例如编程助手、聊天机器人等)走进人们的生活，并得到了广泛的使用。各大基础设施供应商随之也开始了为LLM应用提供更好的推理服务的竞争。
LLM推理中使用了大量的加速卡设备，处理一个LLM请求的成本将超过普通的基于关键字查询请求的10倍。因此，增加推理系统的吞吐量，从而减少cost per request变得非常的重要。

二.挑战

当前LLM的核心都是基于自回归的Transformer模型。这个模型的特点是每次基于用户的输入和之前模型的生成结果来生成一个新的token。对于每一个推理请求，模型都会重复上述过程直到生成了代表"截止"的token。这是一个典型的memory-bound(大部分时间耗在等待内存搬运)的过程，这使得GPU的利用率低下，从而显著的限制了整个系统的吞吐。

由上图可知，增大batch size可以显著的提升系统吞吐。影响batch size的主要制约因素就是显存的大小。vLLM由于对显存的合理利用，可以允许较大的batch size。

三.解决方式

以nvidia A100 GPU为例，在LLM推理过程中，有三部分的显存消耗：

1. 大约有65%的显存用来存放LLM的参数，这一部分是常量，在整个推理过程中保持不变。
2. 有大于30%的显存用来存放KV Cache，这是由前面的token生成过程中缓存的，用来生成新token的上下文信息。
3. 剩下一小部分显存用来存放模型推理过程中生成的临时张量数据(包括激活值等)。

1中的常量存储部分无法进一步优化，3中的临时数据存储占比很小，从而可以优化2中的KV cache可以显著的降低显存占用，从而提升系统吞吐。


不同于其他其他类型的深度学习负载的张量，LLM的KV Cache具备两个特征：

• cache会在模型生成token的过程中发生伸缩
• cache的长度和生命周期无法提前预知

这些特征就带来了内存碎片的问题。受到操作系统的虚拟内存页机制的启发，提出了PageAttention机制。
KV cache被分割成了若干个block，每个block中包含了固定数量token的K和V。每个block在显存中不需要连续存储。

PageAttention机制与操作系统内存页管理类比：

• block就像是操作系统中的内存页(page)
• token就像操作系统中待存储的字节
• 推理请求就像操作系统中的进程

PageAttention通过以下方式减少内存碎片：

• 相对较小的block
• 每个block都有相同的显存占用大小，根据需要为每个block分配显存
• 在不同的sequences和请求之间，显存可以按照block粒度进行复用

vLLM实现了基于PageAttention的block显存管理和抢占式的请求调度，收获了接近于0浪费的KV cache管理和2至4倍的吞吐提升。

1. vLLM整体架构

vLLM采用了一个中心化的调度器来协调分布式GPU节点的任务执行。

• KV Cache Manager:
  它按照Scheduler的调度指令，采用PageAttention的方式来管理各个GPU节点上的物理内存。它以固定大小的KV blocks来组织KV Cache，就类似于操作系统中的虚拟内存页。
  一个推理请求的KV cache被表示成一系列的逻辑KV blocks。随着新的token和它们的KV cache的产生，这些逻辑KV blocks按照从左到右的顺序被填充。
  最后一个KV block的没有被填充的位置，为将来的token所预留。
  KV cache manager同时管理着block table(维护对应物理block和逻辑block的映射关系，以及填充位置的数量)
• Block engine
  在GPU节点上的block engine负责分配一段连续的显存，并将它们分割成若干个物理KV blocks。

举个例子：

① "Four score and seven yeas ago our"是一个有7个token的提示词，vLLM将前两个逻辑KV block（Block 0和Block 1）映射到两个物理KV block(Block 7和Block 1)。
在prefill步骤中，通过传统的自注意力算法，vLLM生成提示词和第一个被生成token的KV cache。vLLM将前4个token的KV cache储存在逻辑block 0中，将接下来的3个token存在逻辑block 1中。剩余空间为接下来的自回归生成的词语预留。

②在第一个自回归decoding步骤里，vLLM通过PagedAttention算法在物理block 7和物理block 1上生成新的token。由于最后一个逻辑block还有一个槽空余，新生成的KV cache会被存放在那里，同时block table中的 # filled 字段对应的值被更新(由3变成4，代表着物理block 1有4个位置已经被填充)。

③在第二个decoding步骤中，最后一个逻辑block已经满了，vLLM将新生成的KV cache存放在一个新的逻辑block中。vLLM为它分配了一个新的物理block(物理block 3)，并将它与逻辑block的映射存在block table中。

从全局上来说，对于每个decoding迭代，vLLM首先选择一组待选的sequences组成batch，并为新获得的逻辑block分配物理block。之后，vLLM将所有当前迭代的输入token(例如，所有提示词的token和最近的用作生成下一个词的token)合并成一个sequence，喂给LLM。在LLM的计算过程中，vLLM使用PageAttention kernel来访问之前存储在逻辑KV blocks的KV cache，并把新生成的KV cache存入物理KV block中。
当在一个KV block中存储多个token时，PagedAttention kernel可以并行处理不同位置的KV cache，因此增加了硬件的利用率并减少了延迟。然而过大的block size也会增加显存的碎片程度。

随着更多的token以及它们的KV cache的生成，vLLM动态的为新的逻辑block分配物理block。所有的这些block都被从左到右的填充，只有当之前的所有block都被填充满了，才会分配一个新的物理block。vLLM将一个推理请求的显存浪费限制在了一个block中，从而有效的利用所有显存。

一旦一个推理请求完成了它的生成过程，它的KV block空间就可以被释放，从而用来存储其他请求的KV cache。

上面的图展示了一个vLLM管理两个sequences的显存的例子。两个sequences的逻辑block被映射到了不同的物理block，这些物理block由GPU节点上的block engine所预留。
所有sequence的相邻的逻辑block不需要被分配在连续的物理GPU内存空间中。并且这些物理block的空间可以被所有sequences有效的利用(例如Request A和Request B的某些逻辑block可以被映射成相同的物理block)。

2. 其他decoding场景的应用

2.1 并行采样(Parallel sampling)

LLM基于一条提示词生成多个输出，用户可以从多个输出中选择一个。在parallel sampling中，一个请求产生的多个输出可以通过PagedAttention共享输入的这条提示词的KV cache。

根据下图来描述一下这个过程：

这种图展示了两个输出的parallel decoding过程。这两个输出(Sample A, Sample B)共享一个提示词，所以在提示词阶段，只为提示词预留一份存储空间即可：逻辑block 0和逻辑block 1都被映射到了物理block 7和物理block 1上。除此之外，我们为每个物理block增加了一个引用计数(reference count)。在图中的例子里，物理block 7和物理block 1的引用计数值都是2(因为同时被Sample A1和Sample
A2的逻辑block引用了)。
在LLM的生成阶段，两个不同的输出就需要不同的空间来存储KV cache了。vLLM实现了基于需要被修改的block粒度的copy-on-write机制，这与操作系统虚拟内存中实现的copy-on-write非常类似(每当fork出一个新进程时触发CoW)。
还是使用上图来举例解释一下vLLM中的copy-on-write机制：

• 当A1需要写入它的最后一个逻辑block时(逻辑block 1中的黄色部分)，vLLM识别到这个逻辑block所对应的物理block(物理block 1)的引用计数大于1。
• 此时vLLM将新分配一块物理block(物理block 3)，之后通知block engine把物理block 1中的数据复制到新分配的物理block 3，然后将物理block 1的引用计数减为1。
• 接下来当A2写入物理block 1时，它的引用计数已经变为1了，A2就可以直接将它新生成的KV cache写入到物理block 1了。

2.2 Beam search

2.3 Shared prefix

2.4 总结

vLLM之所以能够支持多样化的内存共享和访存模式，是因为vLLM将不同sequences之间的复杂的内存共享机制，通过一个公共映射层隐藏了起来。LLM和它的执行kernel只能看到一组物理block id，不需要再去处理不同sequences之间的共享模式了。

3. 调度和抢占

当推理请求量超过系统的容量负荷时，vLLM需要评估请求的优先级。vLLM对所有请求都采用FCFS(first-come-first-server)调度策略来保证公平和防止饥饿。当需要抢占一些请求时，它确保最早到达的请求被响应，最新到达的请求被抢占。

LLM的输入的长度无法预知，它的输出的长度由输入的提示词和模型本身决定，同样无法预测。
随着请求和模型的输出的增长，vLLM的物理block将被耗尽，从而无法储存新的KV cache。这就需要在适当的时候，对已有的block进行驱逐。

• 应该驱逐哪些block？
  因为一个sequence的所有block会被同时处理，所以驱逐的策略也是all-or-nothing，即对某个sequence的策略进行全部驱逐或全部保留。
• 被驱逐的block又需要再被使用，如何恢复？

  1. Swapping
     当物理block被耗尽时，vLLM将选择一组block，将它们从显存拷贝到内存。vLLM架构图中的CPU Block Allocator来管理被交换到内存的物理block。当整个驱逐过程完成之后，vLLM才会继续接受新的请求。当请求结束，其对应的block会被释放，这时之前被驱逐的block会被再换回显存。在这个设计中，被换出到内存的block空间，不会超过当前KV cache的显存大小。

  2. 重新计算
     重新计算被抢占的block的KV cache。这个计算过程会比原始的计算延迟高(as the tokens generated at decoding can be concatenated with the original user prompt as a new prompt—their KV cache at all positions can be generated in one prompt phase iteration.)

Swapping和重新计算这两种如何选择，取决于内存与显存之间的带宽，GPU的算力等因素。

4. 分布式执行

当模型大小超过一个GPU的显存时，就需要将模型分割到不同的GPU设备进行分布式推理。vLLM支持Megatron-LM的张量并行策略。
attention操作在attention head维度被切分，每个SPMD(single program multiple data)处理器会处理一部分attention head。
即使是模型并行执行，每个模型分片依然会共享相同的输入tokens，所以需要用到KV cache。因此，vLLM在中心化的scheduler中提供了一个KV cache manager(见架构图)。不同的GPU worker共享这个KV cache manager(kv cache manager维护着逻辑block与物理block的映射关系)。每个GPU worker取得由scheduler所提供的物理block ID，然后进行计算。虽然每个GPU worker共享了物理block ID，但是一个worker只存储它所负责的attention head的那一部分KV cache。

文章详情