LLM推理优化笔记：vLLM原理与PagedAttention详解
LLM推理优化笔记：vLLM原理与PagedAttention详解
在大语言模型（LLM）的推理过程中，如何高效管理内存是一个关键挑战。本文将详细介绍vLLM团队提出的PagedAttention算法，以及基于该算法构建的vLLM推理系统。通过借鉴操作系统中的分页机制，PagedAttention能够更有效地管理KV缓存，从而显著提升LLM推理的效率和资源利用率。
LLM推理过程中的内存挑战
在LLM进行推理时，每个请求的KV缓存内存占用非常大。例如，PagedAttention论文中的图1显示，一个13B参数规模的LLM在推理时，KV缓存存储占到了30%以上的显存。而且，随着推理过程的进行，KV缓存会动态增加或减小。如果KV缓存管理不当，会导致内存碎片化和冗余KV缓存，造成大量的内存浪费（如论文图2所示）。
LLM推理过程概述
给定一个prompt，LLM推理服务的生成过程分为两步：
prefill阶段（也称为prompt阶段）：将prompt的tokens作为输入，计算第一个新token的概率。在这个过程中会计算出key向量和value向量。由于全部prompt tokens是已知的，这个阶段可以通过矩阵乘法操作进行并行计算，因此是compute-bound的。
decoding阶段（也称为自回归生成阶段）：按顺序一个接一个地生成新的token，直到生成序列达到最大长度或遇到终止符。这个阶段的计算因为每一次计算都依赖前一次迭代的结果，无法并行计算只能使用matrix-vector multiplication，是memory-bound的。
PagedAttention算法原理
受操作系统中分页机制的启发，PagedAttention允许在非连续内存空间存储连续的key和value向量。具体来说，PagedAttention将每一个序列的KV缓存划分为KV blocks，每个block包含固定数目的token数（默认block size为16）。通过block-wise计算来完成attention计算，如下式所示：
vLLM系统架构
基于PagedAttention的vLLM架构如下图所示，Scheduler协调分布式GPU worker的执行，KV cache manager基于PagedAttention来管理KV缓存，其通过scheduler的指令来管理GPU worker上的物理KV缓存内存。
KV Cache Manager
按照PagedAttention的思想，将KV缓存划分为固定尺寸的KV blocks。将KV blocks分为logical KV blocks和physical KV blocks。将一个请求的KV缓存表示为一系列的logical KV blocks，它是连续的，按token顺序从左到右填充KV block，只有最后一个KV block的空闲位置是保留给将来生成的token使用的。
block manager将GPU worker上的GPU DRAM或者CPU RAM划分为physical KV blocks。维护block tables来记录每个请求的logical KV blocks和physical KV blocks之间的映射，block table的每一行记录logical block对应的physical blocks以及填充位置个数。
并行采样与Beam Search
有时候对于单个prompt输入，LLM需要生成多个采样输出。在这种场景下，一个请求对应的多个输出共享相同的prompt，vLLM通过其PageAttention和分页内存管理，可以很容易实现输入prompt的KV缓存共享从而节省内存。
Beam search是在机器翻译等常用的解码方法，它依赖于beam width参数k，用来定义在每一步时top k候选词。vLLM管理k=4的beam search的KV blocks，在图中的虚线之前，每一个候选序列有4个logical blocks，所有序列都共享第一个block 0，Beam candidate 3在第二个block开始与其他候选序列不一样，Beam candidate 0-2共享前3个block在第四个block开始不一致。
共享前缀与调度策略
一般在使用LLM时，用户会使用包含了任务指令、输入输出样例等的system prompt，system prompt会作为最终用户输入给LLM的prompt的一部分，也就意味着许多请求都会共享一个前缀。在vLLM中会保留预先定义的共享前缀对应的physical blocks，这样当用户请求的prompt包含共享前缀时只需要将其logical blocks映射到已经缓存好的physical blocks就可以了（最后一个block需要copy-on-write机制），可以省略掉共享前缀部分的计算。
对于内存管理，vLLM执行的是all-or-nothing驱逐策略，也就是对一个待驱逐请求的所有blocks都进行驱逐。同时因为vLLM对于一个请求会生成多个输出序列(比如beam search请求的beam candidates)一起作为一个序列组(sequence group)来调度，因为这些序列会共享内存，所以一个序列组的序列总是被同时preempted或者重新调度。
对于被驱逐的序列，vLLM有两种技术来实现恢复：Swapping和Recomputation。实验表明对于较小的block size, Recomputing速度更快，对于较大的block size，Swapping更快，而中等大小的block size(16-64)两者速度差不多。
分布式执行
vLLM通过支持Metatron-LM风格的模型并行来支持分布式设置。如前面vLLM架构图所示意，它只有一个中央调度器，不同GPU worker共享KV缓存管理器。