
隐藏的内存消耗:KV缓存
大模型的核心在于Transformer模型,它每次生成一个单词(token)。为了高效输出,模型需要记住来自之前单词的“上下文”,并将其存储在KV(键值)缓存当中,借此实现对话的短期记忆。
问题在于,这套KV缓存极其庞大,且每条请求的大小都会动态变化。现有系统多将KV缓存存储在连续的单一内存块中,因此往往带来两个主要问题:
1. 内存碎片
内部碎片
系统会为每条请求预先分配一大块内存,并假设最大的可能输出长度(如2048个token)。但如果一条请求只生成一条短输出,则大部分预留内存都会闲置,因此造成严重浪费。
外部碎片
由于不同请求会预留大小不一的内存块,GPU内存中会分散大量无法使用的小块。因此即使总可用内存充足,也难以容纳新请求。从观察数据来看,现有系统中只有20.4%至38.2%的KV缓存内存被实际用于存放token状态,其余部分被白白浪费。
2. 无法共享内存
并行采样或波束搜索等高级解码技术常会依据单条提示词生成多条输出,因此可以共享部分KV缓存。但现有系统无法轻松共享这部分内存,因为各序列的KV缓存位于各自独立的连续块当中。
这种低效现状严重限制了系统能够同时处理的请求数量,直接影响到系统吞吐量(每秒token/请求处理量)。
顿悟时刻:受操作系统启发的PagedAttention
为了解决这些难题,研究人员开发出PagedAttention,其灵感来自操作系统中的一项经典技术:虚拟内存与分页。
其中:
KV块类似页。PagedAttention不使用连续内存,而是将每个序列的KV缓存划分成固定体量且较小的键值块,每个块保存一定数量的tokne键和值。
Token类似于字节,KV缓存中的单一token对应页中的字节。
请求类似于进程。每条大模型请求都像进程一样接受管理,其“逻辑”KV块映射至GPU内存中的“物理”KV块。
PagedAttention如何化解内存难题?
几乎消除碎片问题
由于KV块在物理内存中不再连续,PagedAttention可以按需动态分配各块。由于只需在必要时分配内存,因此消除了内部碎片问题;另外所有块均大小相同,因此消除了外部碎片问题。
灵活的内存共享效果
PagedAttention支持在不同序列间共享KV块,甚至跨不同请求进行共享。例如,在并行采样或波束搜索中,多条输出可以共享初始提示词的KV缓存,从而节省大量内存。它甚至可以对需要由不同序列进行修改的块执行写时复制机制(同样属于操作系统概念),确保高效共享并避免不必要的复制。
vLLM:高吞吐量引擎
构建于PagedAttention之上的vLLM是一款专为高吞吐量目标设计的大模型服务系统。它采用块级内存管理,并配合与PagedAttention协同的复杂调度程序。
vLLM的关键优势包括:
KV缓存内存几乎零浪费。
可在请求内部及跨请求间灵活共享KV缓存。
与FasterTransformer和Orca等最先进的系统相比,vLLM将主流大模型的吞吐量提高了2至4倍,且不会增加延迟。对于较长序列、更大的模型和更复杂的解码算法,这种提升会更加显著。例如,在支持一套包含13B参数的大模型时,vLLM的并发处理请求数甚至比Orca的“预言机”版本(假设完全了解输出长度)高出2.2倍,比Orca(Max)高出4.3倍。它的内存节约效果在并行采样(6.1%-9.8%)和波束搜索(37.6%-55.2%)中同样明显。
大模型支撑的新未来
凭借对操作系统设计经验的巧妙借鉴,PagedAttention显著提高了大模型服务效率,在轻松降低云服务成本的同时,为用户提供更快、响应更加灵敏的大模型应用。它改变了游戏规则,以破除瓶颈的方式更好地支撑起下一代大模型服务。
原文标题:Unlocking LLM superpowers: How PagedAttention helps the memory maze,作者:Raul Leite