在上一篇博客《大模型推理–KV Cache》中详细介绍了大模型推理的decoding阶段可以采用KV Cache来优化重复计算的原理。虽然KV Cache大幅提升了大模型token生成的速度，但是也引入了新的问题，主要有两个：1. KV Cache在长上下文的情况下占用量非常大，导致batch很小，进而影响吞吐量，甚至根本无法支持长上下文；2. 大模型推理的时候无法预知会产生多少token，所以无法给KV Cache预分配空间，现在的通用做法是按照生成token的上限来分配空间，这产生了非常大的空间浪费。问题1是当前大模型推理研究的热点，即通过各种手段降低KV Cache的显存占用情况，我会在后续的博客中给大家介绍这个主题。本博客要介绍的PagedAttention(from: Efficient Memory Management for Large Language Model Serving with PagedAttention)对问题2进行了完全的解决，通过优化空间分配大幅降低了大模型decoding阶段的显存占用，进而可以增大batch，使得吞吐量提升2~4倍，本博客的内容主要是基于该篇论文展开的，论文基于PagedAttention提出了新的推理框架vLLM，目前也被大家广泛使用。
1. 传统推理框架存在的问题
不仔细分析一下大家可能对KV Cache的大小没有直观感受。在PagedAttention论文中给出了一个例子：规模为13B的OPT大模型，单个token的KV Cache就可以达到800KB，计算公式为：2(key和value两个向量) x 5120(embedding维度) x 40(模型层数) x 2(FP16)，当输出token的长度上限为2K时就会产生1.6G的显存占用，而且这还没有考虑batch维度。当然，此处我们只考虑了输出token对应的动态显存占用，实际上输入token(也即prompt)部分也会占据很大显存，这和大模型支持的上下文长度相关，假设大模型支持的上下文为128K，则prompt部分的KV Cache就会有102G的大小，所以才会存在KV Cache压缩的问题，本博客不涉及这个主题。
理想情况下，在整个大模型推理生命周期内，模型参数会占用一部分空间，prefilling阶段会生成prompt的KV Cache，后续decoding阶段逐token生成时KV Cache也会慢慢增长，推理结束所有的KV Cache就被释放，只留模型参数在显存中，如图1所示。虽然输入token的KV Cache我们可以预先知道大小，但是输出token的KV Cache则无法提前获取大小，这就导致很难实现图1所示的理想情况。
在vLLM之前的推理框架在执行decoding阶段时会按照最大输出token长度分配显存空间，假如最大输出token长度为2K，则除了输入token部分分配的固定显存外，它还会为这个请求额外再分配1.6G的显存用来存储输出token的KV Cache。对应到图1，就是在decoding一开始会有一个明显的显存增长，后续整个decoding阶段的显存占用保持稳定不变。但是在真实场景下，输出token的个数也就几十上百个，五百个顶天了，这就会使分配的大段显存是未使用状态，从而引入了很多内部和外部显存碎片，导致可以并行处理的batch因为显存利用率不高的问题而无法很大。论文中给出了一个Orca推理框架和vLLM显存内部和外部碎片占用比例的对比，如图2所示。
Orca的max模式表示按照最大输出token个数来分配显存，这种情况下就会存在大量的内部碎片(Internal frag.)，内部碎片表示被分配但是未被使用的空间。Pow2模式表示当分配的KV Cache空间不足时，就把分配的空间x2，思路类似于STL中的vector空间分配模式，这种模式下会产生大量外部碎片(External frag.)，外部碎片就是可以被分配的显存，但是因为空间过小而无法被使用。Oracle模式是一种理想模式，预先知道大模型的输出token个数而分配正好的空间，此时不会有内部碎片，但是还是会有相当量的外部碎片。Orca在上述三种模式下都会存在比较大量的预留空间(Reservation)浪费，这个稍微有点难理解，这块空间最终的表现就是给输出token填充KV Cache，但是它填充完得花时间，这块空间被分配之后就只能给当前的请求使用，无法给其他急需空间的请求使用，导致利用率偏低。总体上来说，vLLM之前的大模型推理框架对显存的利用不够充分，图2中的真实利用率只有20.4%到38.2%(38.2%这种理想情况不会出现)，浪费的显存使得可以开启的batch数很小，从而导致吞吐量不够高。
vLLM通过引入PagedAttention机制，将显存的分配力度缩小到操作系统中页的力度，使得显存的浪费大大降低，每个请求分配的显存96.3%都可以被真实的KV Cache填充，大大提升了显存的利用率，进而可以在相同延迟下启动更大的batch，提升大模型的吞吐量。此外，PagedAttention也很容易实现KV Cache的共享，这在之前的推理框架中是很难实现的。
2. PagedAttention
我个人感觉PagedAttention这个名字有点凑热度的嫌疑，叫Paged KV Cache可能会更合理一点，因为它的目的就是要优化KV Cache的显存利用率。还有一点，PagedAttention从操作系统的虚拟内存出发来展开论文总感觉也是想拔高论文，但实际上从内存池的角度来讲可能会更容易理解。我们也不去纠结那么多，就按照论文中的内容去介绍PagedAttention，反正最终目的是希望大家能明白它的原理。
2.1 操作系统中的虚拟内存
不知道大家有没有这样一个疑问：操作系统是如何在只有固定内存的情况下启动这么多进程的，是如何解决地址冲突的。这个问题在很长时间内都困扰着我，后来在做听歌识曲项目中接触到了大页内存，通过学习大页内存我才真正理解了操作系统虚拟内存的概念。
操作系统为了能同时运行多个进程，会为每个进程提供一个虚拟的进程空间，在32位操作系统上，进程空间大小为4G，64位系统为2^64（实际可能小于这个值）。在很长一段时间内，我对此都非常疑惑，这样不就会导致多个进程访存的冲突吗。比如，两个进程都去访问地址0x00000010。事实上，每个进程的进程空间都是虚拟的，这和物理地址还不一样。两个进程访问相同的虚拟地址，但是转换到物理地址之后是不同的。这个转换就通过页表来实现，涉及的知识是操作系统的分页存储管理。
分页存储管理将进程的虚拟地址空间，分成若干个页，并为各页加以编号。相应地，物理内存空间也分成若干个块，同样加以编号。页和块的大小相同。假设每一页的大小是4K，则32位系统中分页地址结构为：
为了保证进程能在内存中找到虚拟页对应的实际物理块，需要为每个进程维护一个映射表，即页表。页表记录了每一个虚拟页在内存中对应的物理块号，如图3。在配置好了页表后，进程执行时，通过查找该表，即可找到每页在内存中的物理块号。
在操作系统中设置有一个页表寄存器，其中存放了页表在内存的始址和页表的长度。进程未执行时，页表的始址和页表长度放在本进程的PCB中；当调度程序调度该进程时，才将这两个数据装入页表寄存器。
当进程要访问某个虚拟地址中的数据时，分页地址变换机构会自动地将有效地址（相对地址）分为页号和页内地址两部分，再以页号为索引去检索页表，查找操作由硬件执行。若给定的页号没有超出页表长度，则将页表始址与页号和页表项长度的乘积相加，得到该表项在页表中的位置，于是可以从中得到该页的物理块地址，将之装入物理地址寄存器中。与此同时，再将有效地址寄存器中的页内地址送入物理地址寄存器的块内地址字段中。这样便完成了从虚拟地址到物理地址的变换。
由于页表是存放在内存中的，这使CPU在每存取一个数据时，都要两次访问内存。第一次是访问内存中的页表，从中找到指定页的物理块号，再将块号与页内偏移拼接，以形成物理地址。第二次访问内存时，才是从第一次所得地址中获得所需数据。因此，采用这种方式将使计算机的处理速度降低近1/2。为了提高地址变换速度，可在地址变换机构中，增设一个具有并行查找能力的特殊高速缓存，也即快表（TLB），用以存放当前访问的那些页表项。由于成本的关系，快表不可能做得很大，通常只存放16~512个页表项。
2.2 PagedAttention对显存的虚拟化
vLLM中将操作系统虚拟内存的概念引入到KV Cache的优化中，利用PagedAttention实现了对显存的虚拟化分页。vLLM和操作系统虚拟内存之间有一个一一对应的关系，vLLM就对应操作系统。操作系统中的页(page)就相当于vLLM中的KV blocks，操作系统中的页一般是4KB，KV blocks一般会设置固定存放8个或16个token的KV Cache。vLLM中正在处理的一个个请求就类似于操作系统中的进程，每个请求申请KV blocks就类似于进程去申请内存页。操作系统利用页表实现虚拟地址到物理地址的转换，vLLM中也有一个类似的概念叫block table，它实现了逻辑KV blocks到物理KV blocks的转换。论文中有一个逻辑KV blocks到物理KV blocks的转换图，如图4所示，和图3非常类似。
如果让我们总结一下PagedAttention的核心点，就是它将对大块显存的占用分割成小块的block来实现，我们无需一开始就分配大块显存存放未知大小的KV Cache，我们按需分配需要多少分配多少，将浪费控制在一个block之内。每个请求的逻辑KV blocks是连续的，但是物理上的KV blocks则不要求连续。
2.3 不同解码策略下的KV Cache共享
有了PagedAttention，vLLM可以很容易实现不同解码策略下的KV Cache共享。这一块的内容在我们平时使用大模型时较少用到，所以我就简单描述一下。除了我们平时用的greedy decoding之外，我们还经常会用到parallel sampling，beam search等解码手段，这两种方法都会有相同的prompt，所以它们在prefilling阶段生成的KV Cache是相同的，利用PagedAttention不同的逻辑KV block可以访问相同的物理KV block，实现KV Cache的共享。Beam search会更进一步，在decoding阶段也会有几率存在KV Cache相同的情况，都可以通过PagedAttention实现KV Cache的共享，这在之前的推理框架里面是很难实现共享的。当然，为了保证访问的正确性，这里还会引入引用计数(ref_count)和写时拷贝(copy-on-write)等概念，大家感兴趣可以去直接读论文。
另一个实现KV Cache共享的场景就是system prompt，这个在实际的应用中经常会出现，就是给大模型输入很多相同的指令式提示，引导大模型按照要求完成推理，比如“你是一个无所不能的程序员，帮我完成代码的翻译”。这种情况下，所有请求的system prompt部分对应的KV Cache是相同的，我们也可以很容易地利用PagedAttention实现共享。
2.4 vLLM的调度与抢占策略
关于PagedAttention还有一个有趣的点需要介绍。之前的推理引擎是预分配大块显存来响应请求，显存不足就不会创建新的请求，与此不同vLLM在采用PagedAttention情况下可能启动了很多推理，推着推着发现物理KV block没有了，这种情况下就需要我们对当前正在处理的某个请求进行抢占(preempt)。vLLM提供了两种策略，第一种就是把最后一个请求的所有KV blocks换出到CPU内存中，这个策略叫swap；另外一个策略就是把最后一个请求的所有KV block都丢弃，等有空闲物理block时再重新计算(recomputation)。第二种策略看似会导致该请求的推理耗时大大增加，但实际上会小于你的预期。这是因为一般情况下某个请求在丢弃的时候已经执行了多次decoding阶段生成了多个输出token，我们在重新计算时可以将prompt+生成的多个token一次性丢进prefilling阶段生成KV Cache，从而把之前执行的多次decoding阶段给省略了，进而加速了重计算的速度。
2.5 基于内存池的理解
论文是从操作系统角度引入PagedAttention的，但是我们也可以从内存池角度来理解PagedAttention。在高性能编程中，为了避免内存碎片，我们针对需求是一系列固定大小块的程序往往会采用内存池的方式来设计。就是我们预先分配很大一块内存，然后分配成一系列固定大小的块，当程序有需求的时候就从内存池中取一块，用完再放回内存池。PagedAttention在真正使用的时候其实就是在全局内存中分配了超大块的连续显存，然后按照KV block的大小分成很多小块，每个请求在推理的时候都去从这个显存池中去获取空闲块存储KV Cache，推理完了之后再把显存块放回到显存池里，就是这么简单，所以我理解PagedAttention有故意拔高论文的嫌疑。
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
