vLLM加速推理的内部原理详解

见贤思齐

文章目录0.引言1.现状2.技术前提2.1基于Transformer2.2自回归生成2.存在哪些问题2.1KVCache太大2.2复杂的解码算法2.3未知的输入和输出长度3.vLLM方案3.1PagedAttention3.2KVCacheManager3.3其他解码场景3.3.1Parallelsampling3.3.2Beamsearch3.3.3共享前缀4.结论参考0.引言vLLM（VirtualizedLanguageLearningModel）是一种用于自然语言处理（NLP）的模型架构或框架，旨在提高大规模语言模型（如GPT等）的性能和效率。论文看这里：《EfficientMemoryManagementforLargeLanguageModelServingwithPagedAttention》代码看这里：https://github.com/vllm-project/vllm本文结合论文详细介绍下为什么需要vLLM，以及vLLM内部的原理是什么，带来了哪些提升。1.现状在模型推理时，GPU的内存分配如下：其中：（1）Parameters保留权重等参数，是静态的，这部分无法优化；（2）KVCache是Transformer的attention机制引入的中间缓存，这部分下面会详细说明；（3）Others是临时激活函数使用，占用比例较小，优化空间不大。从上面GPU的内存分配来看，KVCache是影响推理吞吐量的瓶颈，如果KVCache管理不好，导致一次推理输出的数量太少，就会导致推理速度降低。2.技术前提下面介绍LLM依赖的技术，了解技术前提才能看清LLM推理时会存在哪些问题。2.1基于Transformer现在的LLM基本上都基于Transformer，如果不是基于Transformer，那么vLLM这套机制恐怕不能生效。Transformer的详细介绍参考：《NLP深入学习：大模型背后的Transformer模型究竟是什么？(一)》《NLP深入学习：大模型背后的Transformer模型究竟是什么？(二)》简单列出关键公式，注意力机制公式如下：Attention(Q,K,V)=softmax(QKTdk)VAttention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})VAttention(Q,K,V)=softmax(dk​​QKT​)V其中：Q=XembeddingWQQ=X_{embedding}W_QQ=Xembedding​WQ​K=XembeddingWKK=X_{embedding}W_KK=Xembedding​WK​V=XembeddingWVV=X_{embedding}W_VV=Xembedding​WV​上面向量的维度情况：WQ∈Rembed_dim∗embed_dimW_Q\inR^{embed\_dim*embed\_dim}WQ​∈Rembed_dim∗embed_dimWK∈Rembed_dim∗embed_dimW_K\inR^{embed\_dim*embed\_dim}WK​∈Rembed_dim∗embed_dimWV∈Rembed_dim∗embed_dimW_V\inR^{embed\_dim*embed\_dim}WV​∈Rembed_dim∗embed_dimXembedding∈Rbatch_size∗seq_len∗embed_dimX_{embedding}\inR^{batch\_size*seq\_len*embed\_dim}Xembedding​∈Rbatch_size∗seq_len∗embed_dim是将语句转为嵌入向量之后的结果。2.2自回归生成大模型一般是基于语言模型，即根据前面的语句预测下一个字的概率：P(x)=P(x1)P(x2∣x1)...P(xn∣x1,...,xn−1)P(x)=P(x_1)P(x_2|x_1)...P(x_n|x_1,...,x_{n-1})P(x)=P(x1​)P(x2​∣x1​)...P(xn​∣x1​,...,xn−1​)实际计算时，假设一条语句是x1x2...xtx_1x_2...x_tx1​x2​...xt​，以下是生成语句的步骤：计算KKK和VVV，将本次的KKK、VVV结果缓存起来，并且预测下一个字是xt+1x_{t+1}xt+1​，模型得到输出x1x2...xtxt+1x_1x_2...x_tx_{t+1}x1​x2​...xt​xt+1​再预测下一个字时，计算KKK和VVV就不需要从头开始计算了，x1x2...xtx_1x_2...x_tx1​x2​...xt​的结果已经算过了，本次只需要计算xt+1x_{t+1}xt+1​的结果，得到K′K'K′、V′V'V′一直重复以上步骤，直到输出达到最大长度，或输出了结束token，如这里的KKK、VVV向量涉及到KVCache的核心，缓存的就是这些向量。2.存在哪些问题2.1KVCache太大LLM服务需要为每个请求维护一个键值（KV）缓存，用于存储模型在生成文本时的上下文信息。随着请求数量的增加，KV缓存的大小迅速增长，占用大量GPU内存。对于13B参数的模型，单个请求的KVCache可能就需要数1.6GB的内存。这限制了同时处理的请求数量，进而限制了系统的吞吐量。2.2复杂的解码算法LLM服务通常提供多种解码算法供用户选择，如贪婪解码、采样解码和束搜索（beamsearch）。这些算法对内存管理的复杂性有不同的影响。例如，在并行采样中，多个输出可以共享相同的输入提示的KV缓存，而在束搜索中，不同候选序列的KV缓存可以部分共享，这要求内存管理系统能够动态调整内存分配。2.3未知的输入和输出长度LLM服务的输入和输出长度是变化的，这要求内存管理系统能够适应不同长度的提示。随着请求的输出长度在解码过程中增长，所需的KV缓存内存也会增加，可能会耗尽用于新请求或现有的内存。现有的LLM服务系统通常采用静态连续内存分配策略，会带来三个方面的内存浪费：（1）预留浪费（reserved）。为每个请求预留最大可能序列长度的内存，然而实际请求的长度可能远小于最大长度；（2）内部碎片（internalfragmentation）。内存分配的低效率还会导致内存碎片，进一步降低内存的可用性；（3）外部碎片（externalfragmentation）。有些内存由于过小，无法使用，这些内存则直接浪费了。3.vLLM方案为了解决这些挑战，vLLM提出了一种新的注意力算法PagedAttention，并构建了一个高效的内存管理系统：KVCacheManager，通过分页技术来管理KVCache，从而提高内存的利用效率，减少内存浪费，并支持更复杂的解码算法。这种方法允许在非连续的物理内存中存储连续的键和值，使得内存管理更加灵活，能够更有效地处理LLM服务中的内存挑战。下面是vLLM的架构：3.1PagedAttentionPagedAttention是一种受操作系统中虚拟内存和分页技术启发的注意力算法。它允许将连续的KKK和VVV向量存储在非连续的内存空间中。这一点与传统的注意力算法不同，后者通常要求KKK和VVV向量在内存中连续存储。算法原理：（1）KV缓存分块：PagedAttention将每个序列的KVCache划分为多个KV块（KVblocks）。每个块包含了固定数量的token的KKK和VVV向量。（2）非连续存储：与传统方法不同，PagedAttention允许这些KV块在物理内存中非连续地存储。这种设计使得内存管理更加灵活。（3）按需分配：PagedAttention按需分配KV块，这减少了内存碎片，并允许跨请求或同一请求内不同序列之间的内存共享。3.2KVCacheManagerKVCacheManager是vLLM系统中的一个核心组件，负责以分页的方式高效管理KVCache。这一管理器的设计灵感来源于操作系统中的虚拟内存管理技术，特别是分页机制。主要功能：（1）逻辑与物理块映射：它维护一个块表（blocktable），记录逻辑KV块与物理KV块之间的映射关系。这种分离允许动态地增长KVCache内存，而无需预先为其所有位置分配内存。（2）动态内存分配：KVCacheManager根据需要动态地分配物理KV块，这消除了现有系统中的大部分内存浪费。（3）内存共享：通过页级内存共享，KVCacheManager支持在不同请求之间共享KVCache，进一步减少内存使用。下面是一个例子说明KVCacheManager和PagedAttention的工作机制：（1）一个Block最大存储4个token。开始输入Fourscoreandsevenyearsagoour一共7个token，因此逻辑KV块Block0填充前4个，对应于物理KV块的Block7；剩下3个token填充逻辑KV块Block1，对应于物理KV块的Block1，但是并没有填满，对应的BlockTable的#filled写入3；（2）当自回归生成下一个tokenfathers后，物理KV块的Block1同时写入，并且对应的BlockTable的#filled由3改写成4；（3）自回归生成下一个tokenbrought后，流程类似，不再赘述。3.3其他解码场景3.3.1Parallelsampling（1）有A1和A2并行采样，A1和A2的逻辑Block0的内容一致，实际上在物理KV块中只会存一份，对应Block7，此时Block7的引用计数是2；（2）当A1和A2的Block1都是yearsagoour时，对应于物理KV块的Block1，此时Block1的引用计数是2；（3）在生成阶段时，A1、A2生成下一个字的结果不一样，A1是fathers而A2是mothers，这样Block1会触发写时复制（copyonwrite），复制一份到Block3，同时Block1填入mothers，Block3填入fathers。Block1的引用由2计数减1，变成1。3.3.2Beamsearch与并行解码不同，束搜索（beamsearch）不仅支持共享初始提示块，还支持不同候选之间的其他块共享，并且随着解码过程的推进，共享模式会动态变化，类似于操作系统中由复合分叉创建的进程树。根据上图，介绍vLLM如何管理Beamsearch的KV块其中束宽k=4。（1）图中，虚线左边是迭代之前，右边是迭代后的结果（2）在虚线所示的迭代之前，每个候选序列已经使用了4个完整的逻辑块。Beamcandidate0是Block0->Block1->Block3->Block5Beamcandidate1是Block0->Block1->Block3->Block6Beamcandidate2是Block0->Block1->Block3->Block7Beamcandidate3是Block0->Block2->Block4->Block8所有束候选共享第一个块，即Block0。从第二个块开始，候选3与其他候选不同。候选0-2共享前3个块，并在第4个块处发散。（3）在虚线所示的迭代之后，所有块（Block9～12）都来自于候选1和2。由于原始候选0和3不再在候选之列，它们的逻辑块被释放，相应物理块的引用计数减少。vLLM释放所有引用计数达到0的物理块（Block2、4、5、8）。然后，vLLM为新候选分配新的物理块（Block9-12）来存储新的KV缓存。（4）现在，所有候选共享Block0、1、3；候选0和1共享Block6，候选2和3进一步共享Block7。3.3.3共享前缀共享前缀是指在多个用户请求中重复出现的文本序列，这些序列通常位于请求的开始部分。在大型语言模型的应用中，如机器翻译或聊天机器人，用户经常提供包含指令和示例输入输出的长描述。这些描述被连接到实际任务输入以形成请求的提示（prompt）。vLLM系统通过以下方式优化共享前缀的处理：（1）预存储KV缓存：vLLM可以预先存储共享前缀的KVCache，这样在处理包含该前缀的请求时，就不需要重新计算这部分的KVCache，从而减少冗余计算。（2）映射逻辑块到物理块：在用户输入提示包含共享前缀时，vLLM可以将提示的逻辑块映射到预先缓存的物理块。这样，只有用户特定任务输入的部分需要在模型中进行处理。（3）提高效率：通过这种方式，vLLM可以更高效地利用内存和计算资源，因为共享前缀的计算只需进行一次，而不是在每个请求中重复进行。共享前缀的应用场景包括：（1）机器翻译：用户可能会提供一个包含翻译指令和几个示例翻译对的长描述。这些描述的开始部分（即共享前缀）在多个请求中是相同的。（2）聊天机器人：在聊天应用中，用户与机器人的交互历史可能会被用作上下文，以便生成更准确的响应。这些交互历史的一部分可能会在多个请求中重复出现。4.结论评估表明，与FasterTransformer和Orca等最先进的系统相比，vLLM在相同的延迟水平下将流行LLM的吞吐量提高了2-4倍。参考[1]https://arxiv.org/pdf/2309.06180[2]https://blog.vllm.ai/2023/06/20/vllm.html欢迎关注本人，我是喜欢搞事的程序猿；一起进步，一起学习；欢迎关注知乎/CSDN：SmallerFL也欢迎关注我的wx公众号（精选高质量文章）：一个比特定乾坤