vLLM 加速推理框架

什么是 vLLM？

一言以蔽之: vLLM是一个开源框架，通过PagedAttention算法有效管理大语言模型的注意力内存，提升了吞吐量和内存使用效率。其特点包括24倍的吞吐提升和3.5倍的TGI性能，无需修改模型结构。文章详细介绍了vLLM的离线推理流程，包括KV Cache优化、内存管理和内存共享机制。

vLLM 的特点和优势：

• 采用了 Paged Attention，可以有效管理 Attention 的 keys、values

• 吞吐量最多可以达到 huggingface 实现的24倍，文本生成推理（TGI）高出3.5倍，并且不需要对模型结构进行任何的改变

背景

• LLM 的推理，最大的瓶颈在于显存。

• 自回归模型的 keys 和 values 通常被称为 KV cache，这些 tensors 会存在 GPU 的显存中，用于生成下一个 token。

• 这些 KV cache 都很大，并且大小是动态变化的，难以预测。已有的系统中，由于显存碎片和过度预留，浪费了60%-80%的显存。

核心机制: Paged Attention

• Paged Attention 允许在不连续的内存空间中存储连续的 keys 和 values。 具体来说，Paged Attention 会将每个序列的 KV cache 划分为块，每个块包含固定数量 tokens 的 keys 和 values。 在注意力计算过程中，PagedAttention 内核有效地识别并获取这些块。

• 分块之后，这些 KV cache 不再需要连续的内存，从而可以像在操作系统的虚拟内存中一样，更灵活地对这些 KV cache 进行管理。

• Paged Attention 对于显存的利用接近理论上的最优值（浪费比例低于4%）。通过对显存进行更好的管理，可以使得单次可以使用更大的 batch size，从而进一步利用 GPU 的并行计算能力。

具体来说，Paged Attention 把每个序列的 KV Cache 进行了分块，每个块包含固定长度的 token，而在计算 Attention 时可以高效地找到并获取那些块

推理流程

所以序列在分块之后，只有最后一个块可能会浪费内存(实际中浪费的内存低于4%)。高效利用内存的好处很明显：系统可以在一个 batch 中同时输入更多的序列，提升GPU的利用率，显著地提升吞吐量 .

Paged Attention 的另外一个好处是高效内存共享。例如，在并行采样的时候，一个 prompt 需要生成多个输出序列。 这种情况下，对于这个 prompt 的计算和内存可以在输出序列之间共享。

通过块表可以自然地实现内存共享。类似进程间共享物理页，在 Paged Attention 中的不同序列通过将逻辑块映射到一 样的物理块上可以实现共享块。为了确保安全共享， Paged Attention 跟踪物理块的引用计数，并实现了 Copy-on-Write 机制。内存共享减少了 55% 内存使用量，大大降低了采样算法的内存开销，同时提升了高达 2.2 倍的吞吐量

不同解码策略下的用法

给模型发送一个请求，希望它对 prompt 做续写，并给出多种不同的回答。我们管这个场景叫 parallel sampling。

1. prefill 阶段： vLLM 拿到 Sample A1和 Sample A2，根据其中的文字内容，为其分配逻辑块和物理块。 其中 A1 和 A2 的 block0 和 block1 是共享的。

2. decoder 阶段：当两个 sample 生成的 token 不一致时，触发 Copy-on-Write，即在物理内存上新开辟一块空间。此时物理块 block1 只和 A2 的逻辑块 block1 映射，将其 ref count 减去 1；物理块 block3 只和 A1 的逻辑块 block1 映射，将 其 ref count 设为1。

   ref count：引用计数，即有多少个逻辑 block 引用该物理 block。