1、概述
PagedAttention是一项内存优化技术，用于高效管理大预言推理过程中的KV Cache（键值缓存）。
核心思想：借鉴操作系统的虚拟内存分页机制。传统KV Cache是连续内存块，存在严重内存碎片和浪费。PagedAttention将KV Cache划分为固定大小的物理页（page），逻辑上连续的注意力上下文可以非连续地存储在多个物理页中，就像操作系统用页表管理虚拟内存一样。
2、传统KV Cache的内存管理
（1）推理两阶段
常规LLM推理会分为Prefill和Decode这2个阶段，需要用到KV Cache技术进行加速。
在Prefill阶段，会把prompt过和
后得到的向量保存在Cache_K和Cache_V中。
在Decode阶段，根据prompt的prefill结果，一个token一个token地生成response，然后逐一把新生成地KV值加到cache中。
- 此时，随着prompt数量变多和序列变长，KV Cache也变大，对gpu显存造成压力。
- 另外，由于输出的序列长度无法预先知道，所以我们很难提前为KV Cache量身定制存储空间。
（2）分配存储空间的常规方式
由于推理所生成的序列长度大小是无法事先预知的，所以大部分框架会按照(batch_size, max_seq_len)这样的固定尺寸，在gpu显存上预先为一条请求开辟一块连续的矩形存储空间。然后，这样的分配方法很容易引起"gpu 显存利用不足"的问题，进而影响模型推理时的吞吐量。
3、PagedAttention原理
（1）虚拟内存核心思想
各个进程间独立开发的做法：
- 给每个进程分配一个虚拟内存。每个进程在开发和运行时，可以假设这个虚拟内存上只有自己在跑，这样它就能大胆操作。
- 虚拟内存负责统一规划代码、数据等如何在物理内存上最终落盘。这个过程对所有进程来说都是透明的，进程无需操心。
（2）虚拟内存的分页管理
假如有2个进程，可以将进程1、进程2想成是两本书。代码分布在书的不同page上。我们希望读哪一页，就加载哪一页，而不是一下把两本书都加载进来。当我们不想读某页时，我们也能根据页码将其清空。
现在，我们希望读进程1和进程2的page1，我们就将其加载到物理内存上。虚拟内存会帮我们做好映射，把来自不同进程的这两页分别加载到物理内存对应位置。
虚拟内存的分页管理技术总结：
- 将物理内存划分为固定大小的块，我们称每一块为页（page）。从物理内存中模拟出来的虚拟内存也按相同的方式做划分
- 对于1个进程，我们不需要静态加载它的全部代码、数据等内容。我们想用哪部分，或者它当前跑到哪部分，我们就动态加载这部分到虚拟内存上，然后由虚拟内存帮我们做物理内存的映射。
- 对于1个进程，虽然它在物理内存上的存储不连续（可能分布在不同的page中），但它在自己的虚拟内存上是连续的。通过模拟连续内存的方式，既解决了物理内存上的碎片问题，也方便了进程的开发和运行。
（3）PagedAttention处理单个请求
整个流程如下：
- 请求request可理解为操作系统中的一个进程
- 逻辑内存(logical KV blocks)可理解为操作系统中的虚拟内存，每个block类比于虚拟内存中的一个page。每个block的大小是固定的，在vLLM中默认大小为16，即可装16个token的KV值
- 块表（block table）可理解为操作系统中的虚拟内存到物理内存的映射表
- 物理内存（physical KV blocks）可理解为操作系统中的物理内存，物理块在gpu显存上，每个block类比于虚拟内存中的一个page
1）Prefill阶段
- 划分逻辑块：vLLM拿到这条prompt，先按照设定好的block大小B（本例中B=4），为prompt划分逻辑块（Logical KV blocks）。由于prompt中有7个token，所以vLLM用2个逻辑块（block 0， block 1）来装它们的KV值。其中，在逻辑块1中目前只装了"years", "ago", "hour"这3个token的KV值，有1个位置是空余的。这个位置就被称为保留位（reservation）
- 划分物理块：划分好逻辑块后，我们就可以将其映射到物理块中去了。物理块是实际存放KV值的地方。我们通过一张block table来记录逻辑块和物理块的映射关系，block table的主要内容包括：
- 逻辑块和物理块的映射关系（physical block number）：例如逻辑块0对应物理块7
- 每个物理块上被填满的槽位（# filled）：例如在prefill阶段，对物理块7，其4个槽位都被填满；对物理块1，其3个槽位被填满。
- 正常计算prompt的KV值，并通过划分好的关系填入物理块中。
2）Decode阶段-生成第1个词
- 使用KV cache计算attention，生成第1个词fathers。不难发现，当我们计算时，我们使用的是逻辑块，即形式上这些token都是连续的。与此同时，vLLM后台会通过block table这个映射关系，帮我们从物理块上获取数据做实际计算。通过这种方式，每个request都会认为自己在一个连续且充足的存储空间上操作，尽管物理上这些数据的存储并不是连续的。
- 基于新生成的词，更新逻辑块、物理块和block table。对于block table，vLLM将它filled字段由3更新至4。
- 分配新的逻辑块和物理块。当fathers更新进去后，逻辑块已装满。所以vLLM将开辟新的逻辑块2，并同时更新对应的block table和物理块。
（4）PagedAttention处理多个请求
4、基于PagedAttention的Parallel Sampling优化
Parallel sampling定义：我们给模型发送一个请求，希望它对prompt做续写，并给出N种不同的回答。
传统的方法：将prompt复制N次后组装成1个batch喂给模型，让它做推理，但是这种方式会产生prompt部分KV cache的重复存储。
假设模型的max_seq_len=2048，传统的KV cache可能在显存中分配两块长度是2048的空间。由于prompt一致，这两块2048的空间中存在大量重复的KV cache。
vllm的做法：
假定我们发给模型1个request，这个request中包含2个prompt/sample，记为Sample A1和Sample A2，这两个prompt完全一致，都为Four score and seven years ago our，我们希望模型对这两个prompt分别做续写任务。
（1）首先，Prefill阶段，vLLM拿到Sample A1和Sample A2，根据其中的文字内容，为其分配逻辑块和物理块。
分配逻辑块：对于A1，vLLM为其分配逻辑块block0和block1；对于A2，vLLM为其分配逻辑块block0和block1。需要注意的是，A1的逻辑块和A2的逻辑块是独立的（尽管它们都叫block0和block1），你可以将A1和A2视作操作系统中两个独立运行的进程。
分配物理块：对于A1和A2，虽然逻辑块独立，但因为它们的文字完全相同，所以可以在物理内存上共享相同的空间。所以A1的逻辑块block0/1分别指向物理块block7/1；A2的逻辑块block0/1分别指向物理块block7/1。我们设每个物理块下映射的逻辑块数量为ref count，所以对物理块block7/1来说，它们的ref count都为2。
（2）然后，进入decode阶段，A1和A2各自做推理，得到第一个token，分别为fathers和mothers。
将生成的token装入逻辑块：对于A1和A2来说，将其生成的token装入各自的逻辑块block1。
触发物理块copy-on-write机制：由于fathers/mothers是两个完全不同的token，因此对物理块block1触发复制机制，即在物理内存上新开辟一块空间。此时物理块block1只和A2的逻辑块block1映射，将其ref count减去1；物理块block3只和A1的逻辑块block1映射，将其ref count设为1。
总结起来，vLLM节省KV cache显存的核心思想是，对于相同数据对应的KV cache，能复用则尽量复用；无法复用时，再考虑开辟新的物理空间。
