专栏名称: 学姐带你玩AI
这里有人工智能前沿信息、算法技术交流、机器学习/深度学习经验分享、AI大赛解析、大厂大咖算法面试分享、人工智能论文技巧、AI环境工具库教程等……学姐带你玩转AI!
目录
相关文章推荐
51好读  ›  专栏  ›  学姐带你玩AI

如何轻松、快速且低成本地部署LLM服务?

学姐带你玩AI  · 公众号  ·  · 2024-10-30 18:30

正文

来源:投稿  作者:175
编辑:学姐

今天带来vLLM的论文Efficient Memory Management for Large Language Model Serving with PagedAttention笔记。作者还发布了一篇技术博客:https://blog.vllm.ai/2023/06/20/vllm.html ,本文的动图来自该博客。

高吞吐量的服务大语言模型需要同时批处理足够多的请求。然而,现有系统由于每个请求的kv-cache内存非常大且是动态扩缩的,当管理不当时,这种内存可能由于碎片化和冗余重复而浪费大量,从而限制了批处理大小。

作者提出了PagedAttention,这是一种受到传统虚拟内存和操作系统中分页技术启发的注意力算法。并在此基础上构建了vLLM,这是一种LLM服务系统,能够实现:

  1. kv-cache内存的几乎零浪费;
  2. 在请求之间和请求内部灵活共享kv-cache,以进一步减少内存使用。

vLLM在与当前最先进系统相比时,能够在保持相同延迟的情况下,将流行LLM的吞吐量提高2-4倍。随着序列变长、模型增大和解码算法复杂度增加,这一改进更加明显。

vLLM的源码开源在https://github.com/vllm-project/vllm。

unset unset 1. 总体介绍 unset unset

运行大语言模型的应用非常昂贵,需要大量的硬件加速器,因此提高LLM服务系统的吞吐量——从而降低每个请求的成本——变得越来越重要。

鉴于这些高成本,提高LLM服务系统的吞吐量——从而降低每个请求的成本——变得越来越重要。

大语言模型的核心是自回归Transformer模型。该模型根据输入和目前已生成的输出标记序列,一次生成一个标记。对于每个请求,这一昂贵的过程会重复,直到模型输出一个终止标记。这种顺序生成过程使得工作负载受限于内存,导致GPU的计算能力未得到充分利用,从而限制了服务的吞吐量。

通过将多个请求批量处理,可以提高吞吐量。然而,为了在一个批次中处理多个请求,每个请求的内存空间需要有效管理。例如,图1左展示了在具有40GB显存的A100 GPU上,一个13B参数的LLM的内存分配。大约65%的内存分配给模型权重,这在服务过程中保持不变。接近30%的内存用于存储请求的动态状态。对于Transformer,这些状态包括与注意力机制相关的键和值张量,通常称为kv-cache。

剩余的小部分内存用于其他数据,包括激活值——在评估LLM时创建的短暂张量。由于模型权重是固定的,而激活仅占用GPU内存的一小部分,因此kv-cache的管理方式在确定最大批量大小方面至关重要。当管理不当时,kv-cache内存可能会显著限制批量大小,从而影响LLM的吞吐量,如图1右所示。

现有的LLM服务系统在高效管理kv-cache内存方面存在不足,这主要是因为它们将请求kv-cache存储在连续的内存空间中,而大多数深度学习框架要求张量以连续内存存储。然而,与传统深度学习工作负载中的张量不同,kv-cache具有独特的特性:它会随着模型生成新标记而动态增长和缩小,其生命周期和长度在先验上并不可知。这些特性使得现有系统的方法在两个方面显著低效:

首先,现有系统存在内部和外部内存碎片问题。为了将请求的kv-cache存储在连续空间中,它们会根据请求的最大长度预先分配一块连续的内存,这可能导致严重的内部碎片,因为请求的实际长度可能远小于其最大长度(例如,图11)。此外,即使实际长度在先验上是已知的,预分配仍然低效:在请求的生命周期内,整个内存块都被保留,其他较短的请求无法利用当前未使用的任何部分。此外,外部内存碎片也可能是显著的,因为预分配的大小对于每个请求可能不同。实际上,从图2中的分析结果显示,在现有系统中,仅有20.4%至38.2%的kv-cache内存被用于存储实际的标记状态。

其次,现有系统无法利用内存共享的机会。LLM服务通常使用先进的解码算法,如并行采样和束搜索,这些算法在每个请求中生成多个输出。在这些情况下,请求由多个序列组成,这些序列可以部分共享它们的kv-cache。然而,由于序列的kv-cache存储在不同的连续空间中,因此现有系统无法实现内存共享。

为了应对上述挑战,作者提出了PagedAttention,一种受到操作系统解决内存碎片和共享问题启发的注意力算法:带分页(paging)的虚拟内存。PagedAttention将请求的kv-cache分割成块(block),每个块可以包含固定数量标记的注意力键和值。在PagedAttention中,kv-cache的块不一定存储在连续空间中。因此,可以像操作系统的虚拟内存一样以更灵活的方式管理kv-cache:可以将块视为页面,标记视为字节,请求视为进程。该设计通过使用相对较小的块和按需分配来减轻内部碎片问题。此外,它通过确保所有块具有相同的大小来消除外部碎片。最后,它使内存共享能够在块的粒度上实现,可以跨与同一请求相关的不同序列,甚至跨不同的请求。

在本篇工作中,作者构建了vLLM,一个基于PagedAttention的高吞吐量分布式LLM服务引擎,能够实现kv-cache内存的近乎零浪费。vLLM采用块级内存管理和与PagedAttention共同设计的抢占式请求调度。vLLM支持流行的LLM,例如GPT、OPT和LLaMA等不同规模的模型,包括那些超出单个GPU内存容量的模型。与最先进的系统相比,vLLM在LLM服务的吞吐量上提升了2-4倍,而对模型准确性的影响几乎为零。这些改进在处理较长序列、较大模型和更复杂的解码算法时更为显著。

总之,贡献包括:

  • 识别服务LLM中的内存分配挑战,并量化其对服务性能的影响。
  • 提出PagedAttention,一种在非连续分页内存中操作kv-cache的注意力算法,受操作系统中的虚拟内存和分页启发。
  • 设计并实现vLLM,一个基于PagedAttention构建的分布式LLM服务引擎。
  • 在各种场景下评估vLLM,并证明其在性能上大幅超越了之前的最先进解决方案,如FasterTransformer和Orca。

unset unset 2. 背景知识 unset unset

2.1 基于Transformer的大语言模型

2.2 LLM服务 & 自回归生成

一旦训练完成,LLM通常作为条件生成服务进行部署。对LLM服务的请求提供一系列输入提示标记( , ⋯   , ),LLM服务根据公式(1)生成一系列输出标记

这里将提示和输出列表的串联称为序列。由于公式(1)中的分解形式,LLM只能逐个采样和生成新标记,每个新标记的生成过程依赖于该序列中的所有先前标记,特别是它们的键和值向量。在这个顺序生成过程中,现有标记的键和值向量通常会被缓存以生成未来的标记,这被称为KV缓存(kv-cache)。

注意,一个标记的KV缓存依赖于它所有的先前标记。这意味着在序列中同一标记出现在不同位置时,其KV缓存将是不同的。

给点一个请求提示,LLM服务中的生成计算可以分解为两个阶段:

2.3 LLM的批量技术

通过批处理多个请求,可以提高LLM服务的计算利用率。由于请求共享相同的模型权重,权重移动的开销在批处理中的请求之间得以摊销,当批量大小足够大时,这一开销可以被计算开销所覆盖。然而,将请求批处理到LLM服务并非易事,主要有两个原因。首先,请求可能在不同的时间到达。一种简单的批处理策略要么使较早的请求等待较晚的请求,要么延迟接收请求,直到较早的请求完成,从而导致显著的排队延迟。其次,请求的输入和输出长度可能差异巨大(见图11)。一种简单的批处理技术会对请求的输入和输出进行填充以使其长度相等,但这样会浪费GPU计算和内存。

为了解决这个问题,人们提出了细粒度批处理机制,如cellular批处理和迭代级调度。与传统的在请求级别工作的方式不同,这些技术在迭代级别操作。在每次迭代后,完成的请求会从批处理中移除,并添加新的请求。因此,新请求可以在等待单次迭代后进行处理,而无需等待整个批次完成。此外,通过特殊的GPU内核,这些技术消除了填充输入和输出的需要。通过减少排队延迟和填充带来的低效,细粒度批处理机制显著提高了LLM服务的吞吐量。

unset unset 3. LLM服务中的内存挑战 unset unset

尽管细粒度批处理减少了计算的浪费并使请求的批处理更加灵活,但可以批处理在一起的请求数量仍然受到GPU内存容量的限制,特别是用于存储KV缓存的空间。换句话说,服务系统的吞吐量受内存限制。克服这种内存限制需要解决内存管理中的以下挑战:

大规模KV缓存。 随着请求数量的增加,KV缓存的大小迅速增长。例如,对于13B参数的OPT模型,单个标记的KV缓存需要800 KB的空间,计算为2(键和值向量)× 5120(隐藏状态大小)× 40(层数)× 2(FP16的字节数)。由于OPT可以生成最长2048个令牌的序列,一个请求所需的KV缓存内存可以高达1.6 GB。当前的GPU内存容量在几十GB。因此,即使将所有可用内存分配给KV缓存,也只能容纳几十个请求。此外,低效的内存管理还会进一步减少批量大小,如图2所示。此外,考虑到当前的趋势,GPU的计算速度增长快于内存容量。例如,从NVIDIA A100到H100,FLOPS增加了超过2倍,但GPU内存仍保持在80GB的最大值。因此内存将成为一个日益显著的瓶颈。

复杂的解码算法。 LLM服务为用户提供多种解码算法可供选择,各自对内存管理的复杂性具有不同的影响。例如,当用户从单个输入提示中请求多个随机样本时,提示部分的KV缓存占总KV缓存内存的12%,可以共享以最小化内存使用。另一方面,自回归生成阶段的KV缓存应保持不共享,因为不同的样本结果及其对于上下文和位置的依赖性。KV缓存的共享程度取决于所采用的具体解码算法。在更复杂的算法,如束搜索中,多个请求束可以共享更大部分(最多可节省55%的内存)的KV缓存,而共享模式会随着解码过程的推进而演变。

调度未知的输入和输出长度。 对LLM服务的请求在输入和输出长度上表现出变异性。这要求内存管理系统能够适应多种提示长度。此外,随着请求的输出长度在解码中增长,其所需的KV缓存内存也会增加,这可能会耗尽可用于处理新请求或持续生成现有提示的可用内存。系统需要做出调度决策,例如删除或交换一些请求的KV缓存,以释放GPU内存。

3.1 现有系统中的内存管理

由于当前深度学习框架中的大多数操作符需要张量存储在连续的内存中,之前的LLM服务系统也将一个请求的KV缓存存储为跨不同位置的连续张量。由于LLM的输出长度不可预测,它们基于请求的最大可能序列长度静态分配一块内存,而不考虑请求的实际输入或最终输出长度。

图3展示了两个请求:请求A的最大可能序列长度为2048,请求B的最大序列长度为512。现有系统中的块预分配方案存在三个主要的内存浪费来源:未来token的保留槽、由于对潜在最大序列长度的过度配置导致的内部碎片,以及来自内存分配器的外部碎片。外部碎片永远不会用于生成的token,因为在处理请求之前就已知。内部碎片虽然也未被使用,但这仅在请求完成采样后才会意识到。它们都是纯粹的内存浪费。尽管保留内存最终会被使用,但在整个请求期间保留这块空间,尤其是在保留空间较大时,会占用本可以用于处理其他请求的空间。在图2中可视化了内存浪费的平均百分比,显示出之前系统中的实际有效内存可能低至20.4%。

unset unset 4. 方法 unset unset

作者开发了一种新的注意力算法——PagedAttention,并构建了一个LLM服务引擎——vLLM,以应对第3节中概述的挑战。vLLM的架构如图4所示。vLLM采用了集中式调度器来协调分布式GPU工作节点的执行。KV缓存管理器通过PagedAttention以分页的方式有效管理KV缓存。具体而言,KV缓存管理器通过中央调度器发送的指令管理GPU工作节点上的物理KV缓存内存。

4.1 PagedAttention

为了解决第3节中的内存挑战,作者引入了PagedAttention,这是一种受到操作系统中经典分页思想启发的注意力算法。与传统的注意力算法不同,PagedAttention允许在不连续的内存空间中存储连续的键和值。具体而言,PagedAttention将每个序列的KV缓存划分为KV块。每个块包含固定数量的token的键和值向量,将其称为KV块大小(B)。

可能通过动图更好理解,用Query向量(for)与Block 0中的4个Key向量相乘,同时与Block1中的4个Key向量相乘和Block2中的2个key向量相乘。

总之,PagedAttention算法允许将KV块存储在非连续的物理内存中,从而在vLLM中实现更灵活的分页内存管理。

4.2 KV缓存管理器

vLLM的内存管理器的核心思想类似于操作系统中的虚拟内存。操作系统将内存分为固定大小的分页(page),并将用户程序的逻辑分页映射到物理分页。连续的逻辑分页可以对应非连续的物理内存分页,使得用户程序可以像访问连续内存一样访问。此外,物理内存空间不需要提前完全保留,使得操作系统可以根据需要动态分配物理分页。

一个请求的KV缓存被表示为一系列逻辑KV块,从左到右填充新生成的标记和它们的KV缓存。最后一个KV块的未填充位置被保留以供未来使用。在GPU工作节点上,一个块引擎会分配一块连续的GPU DRAM,并将其划分为物理KV块(交换swapping也在CPU RAM上进行)。

KV块管理器还维护块表(block table),记录每个请求的逻辑和物理KV块之间的映射关系。每个块表项记录一个逻辑块的相应物理块以及填充位置的数量。通过分离逻辑和物理KV块,vLLM可以动态增长KV缓存内存,而不需要提前为所有位置保留内存,这消除了现有系统中的大部分内存浪费,如图2所示。

4.3 使用PagedAttention和vLLM进行解码

下面通过一个示例(如图6所示)来演示vLLM在单个输入序列的解码过程中如何执行PagedAttention并管理内存:

① 类似于操作系统的虚拟内存,vLLM在初始阶段不需要为可能生成的最大序列长度保留内存。相反,它只保留必要的KV块来容纳在提示计算期间生成的KV缓存。在本例中,提示有7个标记,因此vLLM将第一个和第二个逻辑KV块(0和1)映射到物理KV块(7和1)。在预填阶段,vLLM使用传统的自注意力算法生成提示和第一个输出标记的KV缓存。然后,vLLM将前4个标记的KV缓存存储在逻辑块0中,后3个标记存储在逻辑块1中。剩余的位置将用于后续的自回归生成阶段。

② 在第一个自回归解码步骤中,vLLM使用PagedAttention算法在物理块7和1上生成新的标记。由于后一个逻辑块(Block 1)中还有一个位置可用,新生成的KV缓存将存储在那里,并更新Block Table的填充记录(3变成了4)。

③ 在第二个解码步骤中,由于最后一个逻辑块已满,vLLM将新生成的KV缓存存储在一个新的逻辑块中;vLLM为其分配一个新的物理块(物理块3)并将此映射存储在块表中。

也可以参考下面的动图:

在全局范围内,对于每次解码迭代,vLLM首先选择一组候选序列进行批处理,并为新的逻辑块分配物理块。然后,vLLM将当前迭代的所有输入标记(即提示阶段请求的所有标记和生成阶段请求的最新标记)连接成一个序列,并将其输入到LLM。在LLM的计算过程中,vLLM使用PagedAttention内核访问存储在逻辑KV块形式中的先前KV缓存,并将新生成的KV缓存保存到物理KV块中。在一个KV块中存储多个标记(块大小>1)使得PagedAttention内核能够并行处理更多位置的KV缓存,从而增加硬件利用率并降低延迟。然而,较大的块大小也会增加内存碎片化。

同样,vLLM会随着更多标记和它们的KV缓存的生成分配新的逻辑块,由于所有的块都是从左到右填充的,并且只有在所有先前的块都填满时才会分配一个新的物理块,vLLM将一个请求的所有内存浪费限制在一个块内,这样就可以有效地利用所有内存,如图2所示。这使得更多请求可以在内存中进行批处理,从而提高吞吐量。一旦一个请求完成其生成,其KV块就可以释放出来存储其他请求的KV缓存。在图7中,展示了vLLM为两个序列管理内存的示例。这两个序列的逻辑块被映射到GPU中块引擎保留空间内的不同物理块。这两个序列的相邻逻辑块在物理GPU内存中不需要连续,并且物理块的空间可以被两个序列有效利用。

在 PagedAttention 中, 内存浪费只发生在序列的最后一个块中 。在实践中,这导致了接近最佳的内存使用率,仅浪费了不到 4%。事实证明,这种内存效率的提升是非常有益的:它允许系统将更多序列批处理在一起,提高 GPU 利用率,从而显著提高吞吐量。

4.4 应用到其他解码场景

4.3小节展示了PagedAttention和vLLM如何处理基本的解码算法,比如贪婪解码和采样,这些算法以用户提示作为输入,生成单个输出序列。在许多成功的LLM应用中,LLM服务必须提供更复杂的解码场景,展示出复杂的访问模式和更多内存共享的机会。

并行采样。 在基于LLM的程序助手中,LLM为单个输入提示生成多个样本输出;用户可以从各种候选项中选择一个偏爱的输出。目前已暗含一个请求生成一个单个序列。下面假设更一般情况,即一个请求产生多个序列。在并行采样中,一个请求包括多个共享相同输入提示的样本,允许 输入提示的KV缓存也被共享 。通过其PagedAttention和分页内存管理,vLLM可以轻松实现这种共享并节省内存。

图8展示了两个输出的并行解码示例。左边是样本A1的逻辑块,右边是样本A2的逻辑块,中间是对应的物理块。

由于两个输出共享相同的提示(Four score and seven years ago our),因此只在提示阶段为提示状态保留一个副本的空间;两个序列的逻辑块被映射到相同的物理块:两个序列的逻辑块0和1分别映射到物理块7和1。由于单个物理块可以被映射到多个逻辑块,因此为每个物理块引入了引用计数(reference count)。在这种情况下,物理块7和1的引用计数均为2。

在生成阶段,两个输出分别采样不同的输出标记,并且需要为KV缓存分配独立的存储。vLLM在需要多个序列修改的物理块上实施了一种基于块的写时复制(copy-on-write)机制,类似于操作系统虚拟内存中的写时复制技术(例如,创建进程时)。具体来说,在图8中,当样本A1需要写入其最后一个逻辑块(逻辑块1)时,vLLM识别到对应的物理块(物理块1)的引用计数大于1;此时,它分配了一个新的物理块(物理块3),指示块引擎复制物理块1的信息,并将引用计数减少到1。接下来,当样本A2写入物理块1时,引用计数已减少到1;因此,A2直接将其新生成的KV缓存写入物理块1。

总之,vLLM使得大多数用于存储提示KV缓存的空间可以在多个输出样本之间共享(比如这里的物理块Block7被两个样本共享了),唯一的例外是由写时复制机制管理的最终逻辑块。通过在多个样本之间共享物理块,内存使用量可以显著减少,尤其对于较长的输入提示。

这里也有一个动图。

束搜索。 在诸如机器翻译等LLM任务中,用户期望从LLM输出前k kk个最合适的翻译。束搜索被广泛应用于从LLM解码最可能的输出序列,因为它减轻了完全遍历样本空间的计算复杂性。该算法依赖于束宽参数k kk,这决定了每一步保留的前候选者的数量。在解码过程中,束搜索通过考虑所有可能的标记扩展束中的每个候选序列,使用LLM计算它们各自的概率,并从 个候选者中保留前k个最可能的序列,其中 是词汇表的大小。

与并行解码不同,束搜索不仅在不同候选者之间共享初始提示块,还共享其他块,而且共享模式会随着解码过程的推进而动态变化,类似于操作系统中由复合分叉(compound forks)创建的进程树。图9展示了vLLM如何管理用于束搜索示例的KV块(k = 4)。在虚线所示的迭代之前,每个候选序列使用了4个完整的逻辑块。所有束候选者共享第一个块0(即提示)。候选者3从第二个块起与其他候选者偏离。候选者0到2共享前3个块,并在第四个块发生分歧。在后续迭代中,前4个最可能的候选者都源自候选者1和2。由于原始候选者0和3不再是前候选者,它们的逻辑块被释放,对应物理块的引用计数被减少。vLLM释放所有引用计数达到0的物理块(块2、4、5、8)。然后,vLLM分配新的物理块(块9-12)以存储来自新候选者的新KV缓存。现在,所有候选者共享块0、1、3;候选者0和1共享块6,候选者2和3进一步共享块7。

先前的LLM服务系统需要在束候选者之间频繁进行KV缓存的内存拷贝。例如,在图9所示的情况下,虚线之后,候选者3需要复制候选者2的大部分KV缓存以继续生成。vLLM通过物理块共享显著减少了这种频繁的内存拷贝开销。在vLLM中,不同束候选者的大多数块可以共享。只有当新生成的标记位于一个旧的共享块内时,才会应用写时复制机制,就像在并行解码中一样。这只涉及复制一个数据块。

共享前缀。 通常,LLM用户提供一个(长)任务描述,包括指令和示例输入输出,这也被称为系统提示。该描述与实际任务输入连接在一起,形成请求的提示。LLM根据完整的提示生成输出。图10展示了一个示例。此外,共享前缀可以通过提示工程进一步调整,以提高下游任务的准确性。

对于这类应用, 许多用户提示共享一个前缀 ,因此LLM服务提供者可以提前存储前缀的KV缓存,以减少在前缀上的冗余计算。在vLLM中,这可以通过为一组预定义的共享前缀保留一组物理块来方便地实现,类似于操作系统如何在进程之间处理共享库。具有共享前缀的用户输入提示可以简单地将其逻辑块映射到缓存的物理块(最后一个块标记为写时复制)。提示阶段的计算只需在用户的任务输入上执行。

混合解码方法。 之前讨论的解码方法展示了多样的内存共享和访问模式。然而,vLLM能够同时处理具有不同解码偏好的请求,而现有系统则无法高效地做到这一点。这是因为vLLM通过一个常见的映射层隐蔽了不同序列之间的复杂内存共享,该层将逻辑块转换为物理块。LLM及其执行核心只看到每个序列的物理块ID列表,而不需要处理序列之间的共享模式。与现有系统相比,此方法拓展了具有不同采样需求的请求的批处理机会,最终增加了系统的整体吞吐量。

4.5 调度和抢占

当请求流量超过系统容量时,vLLM必须优先处理一部分请求。在vLLM中,采用先到先服务(first-come-first-serve,FCFS)调度策略,确保公平性并防止饥饿。当vLLM需要抢占请求时,将确保最早到达的请求优先服务,而较新的请求则会被优先抢占。

LLM服务面临着一个独特的挑战:LLM的输入提示在长度上可能会有显著差异,而相应的输出长度是未知的,这取决于输入提示和模型。随着请求及其输出数量的增加,vLLM可能会耗尽GPU的物理块来存储新生成的KV缓存。在这个背景下,vLLM需要回答两个经典问题:(1) 应该驱逐哪些块?(2)如果需要再次使用被驱逐的块,如何恢复它们?

通常,驱逐策略使用启发式方法预测未来最久不会被访问的块,并驱逐该块。由于在我们的情况下,所有序列的块是一起被访问的,因此作者实现了一种全或无(all-or-nothing)的驱逐策略,即要么驱逐一个序列的所有块,要么不驱逐任何块。此外,一个请求中的多个序列(例如,一个束搜索请求中的束候选者)被作为一个序列组进行联合调度。由于这些序列之间可能存在内存共享,序列组内的序列始终是一起被抢占或重新调度的。为了回答第二个问题,即如何恢复被驱逐的块,作者考虑了两种技术:

交换。 这是大多数虚拟内存实现使用的经典技术,它将被驱逐的分页复制到磁盘上的交换空间。在我们的情况下,将被驱逐的块复制到CPU内存。如图4所示,除了GPU块分配器,vLLM还包括一个CPU块分配器,用于管理交换到CPU RAM的物理块。当vLLM没有额外的物理块可用于新token时,它会选择一组序列进行驱逐,并将其KV缓存转移到CPU。一旦它抢占一个序列并驱逐其块,vLLM将停止接受新请求,直到所有被抢占的序列完成。一旦请求完成,其块将从内存中释放,被抢占序列的块将被带回来以继续处理该序列。请注意,使用这种设计,交换到CPU RAM的块数永远不会超过GPU RAM中物理块的总数,因此CPU RAM上的交换空间受到分配给KV缓存的GPU内存的限制。

重计算。 在这种情况下,在重新调度被抢占的序列时简单地重新计算KV缓存。注意,重计算的延迟可能显著低于原始延迟,因为在解码时生成的token可以与原始用户提示连接成一个新的提示——在一个提示阶段迭代中,可以生成所有位置的KV缓存。交换和重计算的性能依赖于CPU RAM和GPU内存之间的带宽以及GPU的计算能力。

4.6 分布式执行

许多大语言模型的参数规模超过了单个GPU的容量。因此,有必要将它们在分布式GPU上进行分区并以模型并行的方式执行。这就需要一个能够处理分布式内存的内存管理器。vLLM在分布式环境中表现出色,支持广泛使用的Megatron-LM风格张量模型并行策略。这一策略遵循SPMD(单程序多数据)执行调度,其中线性层被分为多个部分进行块状矩阵乘法,GPU不断通过allreduce操作同步中间结果。具体而言,注意力操作在注意力头维度上进行拆分,每个SPMD进程负责处理多头注意力中的一部分注意力头。

即使在模型并行执行的情况下,每个模型碎片仍然处理相同的一组输入token,因此需要为相同位置配置KV缓存。因此,vLLM在集中调度器中具有一个单一的KV缓存管理器。如图4所示,不同的GPU工作者共享这个管理器,以及从逻辑块到物理块的映射。这一通用映射允许GPU工作者根据调度器为每个输入请求提供的物理块执行模型。尽管每个GPU worker具有相同的物理块ID,但每个worker仅存储与其对应的注意力头相关的部分KV缓存。

在每一步中,调度器首先为批次中的每个请求准备带有输入token ID的消息,以及每个请求的块表。接下来,调度器将该控制消息广播到GPU worker。然后,GPU工作者开始使用输入token ID执行模型。在注意力层中,GPU worker根据控制消息中的块表读取KV缓存。在执行过程中,GPU worker通过all-reduce通信原语同步中间结果,而无需调度器的协调。最后,GPU worker将本次迭代的采样token发送回调度器。

总之,GPU worker不需要在内存管理上进行同步,因为它们只需在每次解码迭代开始时接收所有内存管理信息以及步骤输入。

unset unset 5. 实现 unset unset

vLLM是一个端到端的服务系统,具有FastAPI前端和基于GPU的推理引擎。前端扩展了OpenAI API接口,使用户能够为每个请求自定义采样参数,例如最大序列长度和束宽度k kk。vLLM引擎由Python代码和C++/CUDA代码组成。使用Python开发与控制相关的组件,包括调度器和块管理器,同时为PagedAttention等关键操作开发自定义CUDA内核。对于模型执行器,使用PyTorch和Transformers实现了诸如GPT、OPT和LLaMA等流行LLM。使用NCCL进行分布式GPU工作者之间的张量通信。

5.1 内核级优化

由于PagedAttention引入的内存访问模式在现有系统中并不高效,因此作者开发了几个优化的GPU内核。

(1)融合重塑和块写入(Fused reshape and block write)。在每个Transformer层中,新的KV缓存被分割成块,重塑为一种优化了块读取的内存布局,然后保存到由块表指定的位置。为了最小化内核启动开销,将它们融合为一个单一的内核。

(2)融合块读取和注意力(Fusing block read and attention)。调整了FasterTransformer中的注意力内核,以根据块表读取KV缓存并实时执行注意力操作。为了确保内存访问的合并,为每个块分配一个GPU warp进行读取。此外,还支持在请求批次内的可变序列长度。

(3)融合块复制(Fused block copy)。由写时复制机制发出的块复制操作可能在不连续的块上进行。如果使用cudaMemcpyAsync API,这可能会导致大量小数据移动的调用。为了降低开销,作者实现了一个内核,将不同块的复制操作批处理为一次内核启动。

5.2 支持多种解码算法

vLLM使用三种关键方法实现各种解码算法:fork、append和free。fork方法从现有序列创建一个新的序列。append方法将新的token附加到序列中。最后,free方法删除该序列。例如,在并行采样中,vLLM使用fork方法从单个输入序列创建多个输出序列。然后,它在每次迭代中使用append方法向这些序列添加新token,并使用free方法删除满足停止条件的序列。vLLM还在束搜索和前缀共享中应用了相同的策略。

unset unset 6. 评估 unset unset

在ShareGPT数据集上,vLLM的请求率可以比Orca(Oracle)高出1.7×–2.7×,比Orca(Max)高出2.7×–8×,同时保持类似的延迟。这是因为vLLM的PagedAttention能够有效管理内存使用,从而能够批处理比Orca更多的请求。例如,如图13a所示,对于OPT-13B,vLLM可以同时处理比Orca(Oracle)多2.2×的请求,以及比Orca(Max)多4.3×的请求。与FasterTransformer相比,vLLM的请求率最高可以高出22倍,因为FasterTransformer没有使用细粒度的调度机制,并且在内存管理上效率低下,类似于Orca(Max)。

图12的第二行和图13b显示了Alpaca数据集上的结果,呈现出与ShareGPT数据集相似的趋势。一个例外是图12(f),在这里vLLM相对于Orca(Oracle)和Orca(Pow2)的优势不那么明显。这是因为OPT-175B的模型和服务器配置(见表1)允许有大容量的GPU内存空间用于存储KV缓存,而Alpaca数据集的序列较短。在这种情况下,尽管它们的内存管理效率不高,Orca(Oracle)和Orca(Pow2)仍然能够批处理大量请求。因此,系统的性能变得以计算为绑定,而不是以内存为绑定。

unset unset 7. 消融研究 unset







请到「今天看啥」查看全文