【vLLM】使用PagedAttention 进行大型语言模型的高效内存管理

重磅推荐专栏：
《大模型AIGC》
《课程大纲》
《知识星球》

本专栏致力于探索和讨论当今最前沿的技术趋势和应用领域，包括但不限于ChatGPT和Stable Diffusion等。我们将深入研究大型模型的开发和应用，以及与之相关的人工智能生成内容（AIGC）技术。通过深入的技术解析和实践经验分享，旨在帮助读者更好地理解和应用这些领域的最新进展

1. 简介与背景

大型语言模型（LLM, Large Language Models）是人工智能领域的一种深度学习模型，它们通过处理大量的文本数据来学习语言的模式，并能完成诸如文本生成、翻译、问答等多种任务。这些模型通常包含数十亿个参数，需要相当大的计算资源来进行训练和服务（即推理或预测）。在服务阶段，模型的参数、中间计算结果（激活值）、以及键值缓存（KV cache）都需要占用显存空间。

在这里插入图片描述
如上图所示，当在一个 NVIDIA A100 GPU 上部署一个拥有130亿参数的大型语言模型时的显存布局情况。这里有几个关键点：

灰色部分表示的是模型参数，这些参数在整个服务期间一直驻留在GPU显存中。
红色部分是指KV缓存，它用于加速推理过程中的重复计算，这部分显存会随着每个请求被分配和释放。
黄色部分表示短暂使用的激活显存，这是在计算过程中临时使用的显存。
右侧的图展示了vLLM如何平滑了现有系统中KV缓存显存的增长曲线，从而显著提高了服务吞吐量。这意味着vLLM能够更有效地管理显存，减少由于显存快速膨胀而导致的服务性能下降。

上图显示了不同LLM服务系统在实验中浪费的平均显存百分比。这表明不同的系统在处理同样的任务时效率可能有很大差异，而某些系统可能会因为设计上的原因导致更多的显存浪费。

当我们在GPU上部署和运行大型语言模型时，需要处理大量的请求。每个请求在处理过程中都会生成一系列的输出词元（tokens），而这些词元的相关信息会被存储在一个称为KV缓存的地方。KV缓存对于加速模型推理过程非常重要，因为它保存了之前所有生成词元的上下文信息，使得模型可以根据这些历史信息更有效地生成新的词元。

然而，现有系统在管理KV缓存方面存在三个主要问题：

预留显存：为了保证每个请求都能有足够空间来存储其可能的最大长度序列，系统会预先分配一大块连续的显存。这意味着即使一个请求实际上只用了很小一部分显存，其他请求也不能使用这块预留的显存，从而造成了浪费。
内部碎片化：由于每个请求的实际序列长度往往比预估的最大长度要短得多，因此预分配的显存中有很大一部分从未被使用过。这种现象被称为内部碎片化，它限制了可以同时处理的请求数量。
外部碎片化：不同请求之间的预分配显存大小不一样，这会导致所谓的外部碎片化。即使总的可用显存足够多，但由于这些不规则的空闲块无法组合成足够的连续空间供新请求使用，结果还是会有许多显存得不到有效利用。
这些问题共同作用，导致了宝贵的GPU显存资源没有得到充分利用，并且限制了系统的吞吐量——即单位时间内能够处理的请求数量。如图3所示，展示了两个不同的请求A和B，它们分别具有不同的最大可能序列长度。通过这种方式，可以看到当前系统中存在的显存浪费情况。

为了解决上述挑战，研究者们提出了像PagedAttention这样的新技术，旨在模仿经典虚拟显存和操作系统中的分页技术来优化KV缓存管理。这种方法可以几乎消除KV缓存显存浪费，并允许在同一时间点灵活地共享KV缓存，进一步减少显存使用，最终提高了服务效率。

2. 大模型服务的显存挑战

然而，现有系统在管理KV缓存方面存在三个主要问题：

预留显存：为了保证每个请求都能有足够空间来存储其可能的最大长度序列，系统会预先分配一大块连续的显存。这意味着即使一个请求实际上只用了很小一部分显存，其他请求也不能使用这块预留的显存，从而造成了浪费。
内部碎片化：由于每个请求的实际序列长度往往比预估的最大长度要短得多，因此预分配的显存中有很大一部分从未被使用过。这种现象被称为内部碎片化，它限制了可以同时处理的请求数量。
外部碎片化：不同请求之间的预分配显存大小不一样，这会导致所谓的外部碎片化。即使总的可用显存足够多，但由于这些不规则的空闲块无法组合成足够的连续空间供新请求使用，结果还是会有许多显存得不到有效利用。

这些问题共同作用，导致了宝贵的GPU显存资源没有得到充分利用，并且限制了系统的吞吐量——即单位时间内能够处理的请求数量。
在这里插入图片描述

如图所示，展示了两个不同的请求A和B，它们分别具有不同的最大可能序列长度。通过这种方式，可以看到当前系统中存在的显存浪费情况。

3. vLLM 的架构与工作原理

在这里插入图片描述

3.1 PagedAttention

为了理解 PagedAttention，我们可以把它想象成一个图书馆的管理系统，只不过这个图书馆存储的是大型语言模型（LLM）推理过程中使用的“书”——即注意力机制中的键（key）和值（value）向量。这些“书”并不是一直放在固定的架子上，而是可以根据需要被动态地分配到不同的“书架”（显存块）中。