KV cache想理解paged attention我们需要先理解KV cache。在现在主流大模型的decode阶段每次新生成一个token需要和前面所有token的KV进行注意力计算由于前面的KV一旦计算出来就不变了不受当前token的影响所以可以缓存下来不用每次都重新计算这就是KV cache。如下图当前输入是token3预测下一个token需要利用token1-3的KV计算注意力朴素做法是重新使用token1-3计算出KV1-3。而利用KV cache的做法是把前两个token的KV都缓存下来图中紫色表示缓存只用计算token3的KV。但这样随之而来的问题是需要给KV cache预分配内存这就引出了系统领域的一个大问题。内存分配在大模型之前写传统程序或者设计操作系统都会遇到这个问题如何高效地给一个进程分配内存。最简单的做法是分配一块一定够用的定长内存我们在C编程里声明一个静态数组以及早期OS分配连续内存都是这么做的。这实现简单但是会带来很多性能问题首先我们需要估计这个进程需要多大的内存估计小了会导致内存不足运行错误估计多了会浪费内存其次分配的内存往往并不是一开始就全用上的在很长的时间内可能都只用了一小部分即使最后都用上了也造成了浪费。开始这段时间这些内存其实可以分配给其他进程的以上两个问题一般编程时还属于比较常见的可以理解。但最隐蔽也最致命的问题是每次都连续分配内存会导致内存碎片这个问题只有站在系统角度才能发现。如下图灰色是未分配的内存对于上面的情况虽然中间的灰色部分加起来也有不少合并到一起也有50K可以分配给一个进程但是他们分散着在连续分配内存的情况下难以利用。熟悉CPP编程的可能会了解对于预分配内存可能浪费的问题在CPP里的解决方法是容器的自动扩容机制开始只申请很小一部分内存当这块内存被利用超过一个比例时申请一块长度翻倍的内存把数据拷贝过去释放原来的内存。这样任何时候使用的内存都只比当前数据量大一点不存在浪费。但这样的问题是重申请内存释放内存以及数据拷贝非常耗时虽然节约了内存但是增加了时间这对于大模型推理这种极度要求性能的场景没有意义。内存分页在大模型之前操作系统的开发者找到的解决方法是内存分页虚拟内存把一个进程需要的内存划分成很多个较小的页进程能看到的页是虚拟的系统保存一个页表把每个虚拟的页映射到物理页。这样即使进程开始预分配很大内存系统也不会真的分配内存给进程而是在进程运行时访问某个具体页时发现没有对应的物理页触发缺页中断 Page Fault才分配一个物理页和这个虚拟页绑定。这样完美解决了前面预分配连续内存的问题不用担心预分配多了浪费启动时分配的只是虚拟内存运行时才按需调用物理内存按页组织只有页内能产生内存碎片页外没有内存碎片。paged attention铺垫了这么多终于到paged attention了。paged attention是VLLM团队借鉴传统系统优化中的内存分页技术提出的KV cache优化方法。KV cache面对的问题比传统系统更严重传统系统中运行的进程需要多少内存开发者往往是可以估算的分配的内存一般都会用上只是在开始不会马上全用上导致一定的浪费。但大模型推理会输出多少个token是难以预测的很多时候只能设置一个输出长度上限然后按这个上限分配KV cache内存这会导致巨大的内存浪费。可能最大长度1024但模型生成了100个token就结束了那么剩下的924个token的内存在推理期间全部被浪费了。对于大模型推理系统来说想要提高效率最重要的就是批处理并行也就是同时计算多个请求的推理系统吞吐量的瓶颈就是批处理大小。而显存资源往往是有限的如果我们给每个请求都预分配很大的KV cache空间能并行的请求就变少了这对追求吞吐率的推理系统来说是硬伤。于是借鉴传统系统的内存分页思想paged attention对序列按token进行分块估计一定数量比如4个token为一块。序列能看到的块是虚拟的经过一个映射表映射到实际的物理内存块这就实现了对KV cache的内存分页和虚拟内存。和操作系统类似这样也可以按需分配消除预分配的浪费问题同时消除了内存碎片。如下图中国 的 历史 非常这前四个token组成第一个块悠久组成第二个块他们在12虚拟块经过页表映射憨批实际映射到1,3物理块对推理序列来说这两块是连续的但在底层是不连续的两块。VLLM的优化不止于此传统OS里页可以共享比如父子进程会共享同一个数据创建子进程是并不会复制一份数据而是会让子进程的虚拟页也指向父进程的物理页只把物理页的引用计数加一只有某个进程要修改共享内存时才在物理内存上复制一份去修改。如下图两个进程在自己虚拟内存上的共享内存经过页表映射后对应的是物理内存中的同一块内存。传统系统中虚拟内存是为了读同一份数据。大模型推理里也有类似的问题多个请求可能某一段token是完全相同的这多见于一个prompt生成多个回答的场景输入token是完全相同的。那对于这些token的KV也完全相同没必要重复计算重复保存可以让一个请求算出来保存到物理块后面的请求的虚拟块都指向这个物理块。只用引用计数自增这个优化叫做sharing KV block。如下图两个翻译请求输入都是一样的只有翻译的输出才开始不一样他们的输入部分对应的前三个块都指向的相同的物理块1,3,4引用计数都是2对于他们第四个块开始翻译了翻译内容不一样才开始指向不同的物理块2,5共享KV块对于采样也更友好在beam search束搜索中束之间会共享前缀这部分前缀可以共享KV块比如下图中三个搜索束都共享中国的历史这个前缀只用把中国的历史 对应的块引用记为3就行。