主要参考:
vLLM核心技术PagedAttention原理
总结一下 vLLM 的要点:
- Transformer decoder 结构推理时需要一个token一个token生成,且每个token需要跟前序所有内容做注意力计算(包括输入的prompt和该token之前生成的token),一般为了避免重复计算,会保存prompt和前序token的 K,V (称为 KV Cache)。但随着大模型时代的到来,输入的 prompt 和生成的回答越来越长,使得 KV Cache 造成的内存压力越来越大。需要对内存使用进行优化。
- 常规的内存分配方式有以下痛点:
- 除了 prompt 部分是完全用满的,reservation frag 在还没有用到的时候是浪费了的(本来可以用来存一些临时变量),internal frag 是完全浪费的(而且batch数据中max_length越大,浪费越多)。至于内存碎片,一般是当多个小操作反复申请和释放内存时,很容易在内存池中形成多个大小不一的碎片。即使最终总内存足够,但如果需要一个较大的连续块,而内存池中只剩下若干零散的块,就会导致内存碎片问题(图中只是举例,真实情况并不是图示的那种每个request产生一个相同大小的内存碎片)
-
采用 PagedAttention虚拟内存的分业管理技术,首先将 KV 分解成多个固定大小的页面(或称为“块”),同时将物理内存也按相同方式划分。这样对于1个进程,我们不需要静态加载它的全部代码、数据等内容。我们想用哪部分,或者它当前跑到哪部分,我们就动态加载这部分到虚拟内存上,然后由虚拟内存帮我们做物理内存的映射。这样通过分块计算,系统只需分配较小的内存块用于当前页面的计算,而不必一次性分配一个超大的连续内存区域。这在内存碎片较多时尤其有利,因为小块内存更容易在碎片化的内存池中找到连续区域。 因此有利于缓解或者说消化内存碎片。而对于 reservation frag 和 internal frag ,在传统方法中,它们属于预先分配的,还没用上的内存。而在分业管理技术中,reservation frag 在还没轮到时,不会预分配。internal frag 则由于不存在预分配,压根就不会存在了。
-
由于分页管理技术是跑到哪,再动态加载这部分到内存,因此可能会出现显存已经打满,但一个batch的计算还没有全部完成的情况。针对这种情况,vLLM采用的是 先到的先算,如果没算完显存就被打满,则把后来的任务中断,并释放其 KV 占用的显存,供先来的任务算完。这里释放后来任务的 KV 占用,既可以是扔掉后面再重新算,也可以是先放到CPU 内存当中,后面再直接调用。因此虽然vLLM看似还是一个batch的数据输入,再返回一个batch的输出,但实际有可能这批数据有的是优先算好,才把资源释放给后面的数据算的。
-
vLLM还采用了 内存复用机制,例如:1.对于同样一个prompt,给出三种不同回答的场景,这个prompt的 KV 就可以在三个数据中都使用,而不用复制三份,即虽然虚拟内存上看是有三份,但这三份在物理内存上都指向一份,实际只消耗了一份内存。2. 在 beam search 场景,top k 的序列往往有很多前置 token 是相同的,这些也可以用类似上面的方法,即虚拟内存各占各的,但相同 token 的在物理内存上指向同一份。
最后,虽然vLLM看似只是通过节省内存提高推理吞吐量(可以用更大的batch),从而达到加速的目的。但除此之外,还有一些方面也加速了推理:
- 传统方法中,经常要分配大块的连续内存,容易因等待内存释放而出现延迟,而vLLM的分业管理技术,降低了大块内存分配的延迟开销
- 避免一次性在多个层次(全局内存、共享内存、寄存器)之间搬运大块数据,提高了内存带宽利用率和整体效率。而且通过分页,vLLM 可以在有限的 GPU 内存内高效管理和访问 KV 数据,确保注意力计算时数据能就近命中,从而提升整体计算速度。