vLLM讲解
AI摘要: 本文探讨了LLM（Language Model）中KV Cache的不足，包括其计算过程中的冗余重复、显存利用率问题以及内存碎片问题。文章指出，尽管通过KV Cache可以缓存中间结果以提升推理速度，但在显存利用率上存在多个不足，如预分配显存未充分利用、显存碎片问题以及无法有效处理不同长度的Prompt导致的显存浪费。接着，文章介绍了操作系统中的Page Attention技术，该技术通过将内存划分为多个page和虚拟内存来优化进程内存分配，为vLLM提供了一种类似解决方案。最后，文章讨论了共享KV blocks的概念，旨在减少显存占用并提高吞吐量。
KV Cache的不足
我们知道，LLM通过计算当前Token和过往所有Token的QK值，并与V向量做加权求和，以此来预测下一个Token。在计算加权求和的过程中，可以发现KV的计算是冗余重复的，因此可以通过KV Cache来缓存中间结果，大大提升推理速度。
然而，KV Cache在显存利用率上存在以下不足：
-
预分配显存，但是利用不到：由于LLM的输出长度并非固定，有的问题回答长，有的问题回答短，如果预分配最大长度的显存，在LLM输出短的回复时造成了大量的浪费
-
预分配显存，但是尚未利用：在一次请求A中，预分配了显存等待KV Cache填充，但是此时另外一个B请求也由LLM生成回答，并且只会回复很短的内容，那么此时却无法利用请求A预分配的显存（反正此时A预分配的显存暂时又用不到）
-
显存碎片问题：即使所有请求生成长度完全一样，但是由于Prompt的长度不同，每次预分配的KV Cache大小也不同，当一个请求生成完毕释放显存，但是下一个请求的Prompt长度大于这次请求的Prompt的长度，因此无法放入刚刚被释放的显存中。
Page Attention
内存碎片问题在操作系统中是常见的问题，不同的进程所需要的内存大小不同，释放时机不同，这和KV Cache遇到的问题如出一辙。操作系统采用了 Page 和 虚拟内存 的方法来解决进程内存分配问题，vLLM便是借鉴了这种技术提出了Page Attention。
在OS中，内存被划分为大小为的多个page，每个进程拥有虚拟申请可以申请不同大小的虚拟内存，虚拟内存和物理内存通过映射表来关联。
类似地，vLLM将显存划分为 KV Block，例如 4个Token的 KV Cache 大小为一个KV Block，每次请求都可以申请虚拟显存，按照KV Block进行分配，虚拟显存和物理显存由映射表关联。这种方式减少了显存碎片，大大提高了显存利用率。最终效果，vLLM一般可以将显存利用率从20%~40%提升到了96%
Shared KV blocks
当多个请求共享一个Prompt，那么为什么不将这个Prompt的KV Block共享呢？这又可以减少显存占用，提高吞吐量。