1.介绍
LLM处理请求的成本比传统的keyword query要高10倍以上,所以如何降低这个成本显得非常重要。
LLMs的核心是一个自回归Transformer模型,这个模型根据输入(prompt)和它到目前为止生成的output token来生成下一个token,直到输出终止token。这种顺序生成的过程使得负载受内存限制,无法充分利用GPU的计算能力,从而限制了服务吞吐量。
上述的情况可以通过批量处理多个请求来提高吞吐量,但是批量处理多个请求需要有效的管理每个请求的内存空间。KV cache是由和attention机制相关的key和value组成的,表示从早期token到按序生成新的token的上下文。
以13B的模型为例,65%是模型参数静态占用,超过30%都是动态的KV cache。由于模型参数恒定的,所以管理KV cache是最重要的优化点。
现有的LLM服务系统不能有效的管理KV Cache:
-
将请求的KV Cache存储在连续的内存中(DL框架要求将张量存储在连续内存中)
-
为了这个目的,需要预先分配一个最大长度的连续内存块<例如2048token>,造成内部严重的碎片化。因为请求实际长度有长有短,动态变化。
-
因为整块在请求的生命周期一直存在,较短的请求就不能使用块中任何当前未使用的部分,造成外部的内存碎片。
-
-
不能利用内存共享的机制
-
模型会使用decoding算法,例如parallel sampling和beam search来产生多个输出;这种多个输出序列其中可以共享部分KV Cache
-
PagedAttention:受操作系统中用于解决内存碎片和共享内存的方案:带分页的虚拟内存
-
将请求的KV Cache划分为块,每个块包含固定数量的token的key和value
-
数据块可以存储在不连续的空间中,将块视为页、token视为bytes、请求视为进程
这样的设计可以使用相对较小的块并按需分配来缓解内存的内部碎片化;同时他也消除了外部碎片,因为块的大小都是相同的;最后,也实现了跨不同序列或请求的共享内存。
vLLM采用了块级内存管理和抢占式调度,适用于GPT,OPT和LLaMA等流行LLMs,提升2-4倍的吞吐量,且不影响准确率。
2.LLMs服务的内存挑战
虽然批处理可以减少计算资源浪费,但是批处理的请求数仍然受到GPU内存容量的限制。克服这种内存限制需要解决内存管理中的以下挑战:
-
Large KV Cache(GPU显存还是重要瓶颈)
-
KV Cache随着请求数量的增加而快速增长,例如13B的OPT,单个token的KV Cache需要800kb的空间。OPT生成可以多达2048个tokens,所以一个请求的KV Cache需要1.6GB,所以大容量显卡容纳的并发请求数也很少。
-
低效的内存管理,会进一步减少batch size
-
-
复杂的decoding算法(beam search)
-
例如一个prompt的输入,产生多个随机采样的请求答案,这里是可以共享KV Cache内存的(12%)
-
-
调度未知的输入和输出长度
-
输入和输出的可变性,要求内存管理系统适应各种长度
-
当输出长度增加,KV Cache内存扩大,需要系统作出调度决策来替换或删除某些请求在GPU的KV Cache
-
现有系统的内存管理
由于当前的深度学习框架的大多数操作要求张量存储在连续内存中,现有的LLM系统服务也将KV Cache存储为不同位置的连续张量。
由于LLM的输出长度不可预测,他们根据请求的最大可能序列长度(max_tokens)来请求静态分配内存块,而忽略了请求的实际输入和输出长度。
下图展示了现有LLM服务系统的内存管理方式:
黄色表示请求A:一个最大可能是2048个token的序列长度;绿色表示请求B:一个最大可能是512个token的序列长度
预分配块内存方案会直接分配一整块连续内存
-
reversed:为未来生成的token预留的内存
-
internal fragmentation:过度预留而实际没有使用的内部内存碎片,直到请求完成才会释放
-
external fragmentation:内存分配器导致的外部碎片(需要找到连续大块内存,会出现外部碎片)
如上图所示,有效内存使用只占了20.4%
PagedAttention执行过程
vLLM在单个输入序列解码过程中执行PagedAttention和内存管理
-
与操作系统的虚拟内存一样,vllm不需要为最初可能生成的最大序列长度分配内存,只保留必要的KV Cache以容纳prompt的计算。
-
如上图所示:逻辑KV blocks(0,1)通过block table映射到物理KV blocks(7和1,GPU DRAM)
-
在预填充阶段,vLLM使用传统的self-attention来生成prompt的KV Cache并存储在对应的物理block
-
-
第一个自回归解码生成token时,使用PA生成新的token,由于逻辑块还剩一个槽位可用,新生成的KV Cache就存在这里,并且更新#filled的记录。
-
当解码第二步,因为最后一个逻辑块也满了,所以vLLM将新生成的KV Cache存储到新的logical block,并将此映射存储块表中。
总体解码迭代而言,每一次都会将之前的所有输入token和新生成的token连成一个序列,并输入到LLM中。