阿里巴巴 & 上海交大提出 DistKV-LLM 分布式 LLM服务系统 | 端到端吞吐性能翻倍，18个数据集上得到验证！

AIGC 先锋科技

发布于 2024-07-08 12:58:53

1460

发布于 2024-07-08 12:58:53

文章被收录于专栏：AIGC 先锋科技

大型语言模型（LLM）的迅速普及已成为推动基于云的LLM服务增长的主要驱动力，如今这些服务对于推动AI应用的发展至关重要。然而，LLM服务的动态自回归性质以及支持极其长上下文长度的需求，要求对大量资源进行灵活分配和释放。这为设计基于云的LLM服务系统带来了巨大的挑战，其中低效的管理可能导致性能下降或资源浪费。为了应对这些挑战，本文介绍了一种新的分布式注意力算法DistAttention，它将KV Cache分割成更小、更易于管理的单元，从而实现分布式处理和存储注意力模块。基于这个想法，作者提出了一种名为DistKV-LLM的分布式LLM服务系统，该系统动态管理KV Cache，并有效协调数据中心中所有可用的GPU和CPU内存。这确保了云端的高性能LLM服务，可以适应广泛的上下文长度。在具有32个NVIDIA A100 GPU（配置为2至32个实例）的云环境中验证了作者的系统，该系统表现出1.03-2.4倍的端到端吞吐量改进，并支持比当前最先进的LLM服务系统长2-19倍的上下文长度，这在18个数据集上的广泛测试中得到证实，这些数据集的上下文长度高达1,900K。

1 Introduction

大型语言模型（LLM）推动了LLM云服务的快速发展，已成为推进AI应用发展的关键基础设施。然而，由于巨大的计算和数据需求，这一发展面临许多挑战。这些服务通常使用多个GPU卡协同工作来处理LLM任务。然而，LLM的动态性质导致了复杂的计算问题。

LLM服务的心脏是自动回归文本生成的内在过程，其中模型一次生成一个单词（或Token）。每个新生成的Token都会附加到现有的文本语料库中，形成LLM内的输入以进行校正。这个过程会持续，直到最终生成单词或Token。至关重要的是，LLM服务的所需内存和计算资源在整个LLM服务中动态振荡，事先无法知道序列的寿命或长度。

自动回归文本生成的动态和迭代性质使得事先规划资源分配变得不可能，在设计云计算上的高效LLM服务系统时，提出了巨大的挑战。特别是在长上下文任务中，扩展的Key-Value（KV）缓存可能在计算实例内超过GPU内存限制，需要立即重新分配资源。这通常涉及启动代价高昂的实时迁移，将任务转移到更强大的实例，或者预分配额外的GPU来处理可能的内存超载。然而，后者可能导致效率低下和资源浪费，尤其是在具有正常上下文长度的任务中。

以前的工作，例如PagedAttention，试图通过促进GPU和CPU内存之间的数据交换（或交换）来解决这些问题。然而，这种方法存在几个限制。首先，PagedAttention的内存交换范围受到单个节点内GPU和CPU内存的限制，因此限制了其容纳极端长上下文长度的能力。其次，尽管其分页策略旨在最小化内存碎片，但它根据请求 Level 交换整个KV Cache，因此无法在分布式云环境中实现更适应和细粒度的调度。最后但同样重要的是，交换出请求的计算中断可能导致运行任务的性能波动，从而危及云服务至关重要的严格服务 Level 协议（SLAs）。

为了解决上述挑战，作者提出了一种新的注意力算法DistAttention，旨在克服这些挑战。DistAttention将KV Cache划分为rBlocks--统一的子块，这些子块有助于支持长上下文长度的LLM服务中的注意力模块的分布式计算和内存管理。与传统的方法不同，DistAttention利用了数据中心中所有可用的GPU或CPU内存资源，特别是那些目前利用率不高的内存。这不仅实现了支持更长的上下文长度，还避免了与数据交换或实时迁移过程通常相关联的性能波动。

在本文中，作者开发了一种名为DistKV-LLM的分布式LLM服务引擎，该引擎与DistAttention无缝集成。DistKV-LLM在管理KV Cache方面表现出色，能够在数据中心中分布式GPU和CPU之间的内存使用方面高效协调。当一个LLM服务实例由于KV Cache的扩展而面临内存短缺时，DistKV-LLM会主动从负担较轻的实例中获取补充内存。

此外，DistKV-LLM引入了一种复杂的协议，该协议可以促进在云中运行的多个LLM服务实例之间高效、可扩展和一致的交互。该协议旨在有效地管理和平衡大量内存资源。此外，DistKV-LLM优先考虑数据局部性和通信优化，这对于解决与KV Cache分布式存储相关的性能挑战至关重要。

总之，作者的工作旨在充分利用数据中心内所有可用的GPU资源，以确保在处理长上下文任务时，LLM云服务的顺利和高效。结合DistAttention和DistKV-LLM，为LLM服务在分布式环境中面临的资源分配和优化挑战提供了一种解决方案。这种方法实现了高效的资源管理，使LLM服务能够有效地处理各种上下文生成任务。

作者在一个配备32个NVIDIA A100 GPU的云环境中对DistKV-LLM进行了全面评估，测试了从2到32个实例的各种分布式系统配置。作者的评估包括18个具有上下文长度达到1,900K的基准数据集。作者的系统在端到端性能方面比最先进的现有工作提高了1.03-2.4倍，并支持2-19倍更长的上下文长度。

本论文做出了以下贡献：

解锁了在数据中心内充分利用所有可用的GPU和CPU内存资源的可能性。同时，它还有效地优化了内存局部性和通信开销，确保了LLM的顺畅和高效云服务。
在一个配备32个NVIDIA A100 GPUs和18个数据集，每个数据集上下文长度可达1,900K的云环境中，作者的系统超过了最先进的工作，支持上下文长度为2-19倍更长，并在标准上下文长度任务中实现了1.4-5.3倍更高的吞吐量。

2 Background

Large Language Models

基于Transformer的大语言模型（LLMs） 已经彻底改变了自然语言处理，提供了从简单的文本生成到复杂的解决问题和对话AI等能力。

2.1.1 Model Architecture

大语言模型（LLMs） 基于Transformer模型的原则构建了一种复杂的架构。例如，GPT-3 是最大规模的模型之一，由多个带有1750亿参数的Transformer块（层）组成，使其能够捕捉复杂的语言模式并生成类似人类文本的能力。一个Transformer块由几个关键组件组成：

QKV线性层 首先将输入传递到Transformer块中。这些层本质上都是全连接的神经网络，将输入映射到三个不同的空间，包括 Query （Q）、Key（K）和Value（V）。

多头自注意力机制，或注意力模块，是Transformer块的核心。它允许模型以不同的方式对待输入序列的不同部分。在多头注意力中，输入经过线性变换多次形成不同的“头”，允许模型在不同的表示子空间中以不同的位置联合关注信息。

前馈神经网络，或FFN模块，通常在自注意力模块之后。这是一个包含两个隐藏层和中间ReLU激活的神经网络，对于不同的位置，该网络是相同的，但每一层的参数不同。

2.1.2 LLM Service

预填充阶段 在推理过程中，LLM首先从用户那里接收提示或输入文本。这个输入文本被处理以理解上下文和所需任务的性质（例如，回答问题、写诗等）。

给定一个包含

个Token

X=[x_{1},x_{2},\ldots,x_{n}]

的提示文本，模型预测下一个Token

x_{n+1}

。表示每个可能Token为

x_{n+1}

的概率分布

P(x_{n+1}|X)

，即

x_{n+1}

被每个可能Token的概率，可以计算为：

P(x_{n+1}|X)=\text{Softmax}(W\cdot h_{n}+b) \tag{1}

在上述公式中，

和

是从Transformer模型最后一层学习的参数，而

h_{n}

是与最后一个Token

x_{n}

Parallelism Method for LLM Service

大语言模型（LLMs） 在服务或推理过程中需要大量的计算能力和内存资源。为了管理这一点，采用了几种并行策略来分配工作负载并加速处理。

2.2.1 Data parallelism

为了处理云计算环境中大量请求，采用了数据并行主义。在数据中心中，将LLMs复制。基础计算单元称为一个_实例_，它具有完整的模型副本。每个实例独立运行，同时处理不同的请求批次。

分组策略。 在每个实例中，分组策略对于提高吞吐量至关重要，可以同时处理更多的请求。由于上下文长度的变化，请求通常具有不同的生命周期，因此需要动态分组策略。由于Transformer、CNN等专有词汇，因此不需要进行翻译。

为了提高在GPU上的LLM服务的吞吐量，已经引入了各种方法，以改善GPU上的LLM服务吞吐量。

2.2.2 Model Parallelism

模型并行主义是一种技术，用于处理无法完全在单个GPU内存中进行推理的LLM。它涉及将模型分跨多个设备或节点。模型并行主义可以主要分为两类：流水线并行主义和张量并行主义。

流水线并行主义。 在流水线并行主义中，模型的层被分片到多个设备上。它涉及将模型分成几个阶段或层，每个阶段都在不同的计算单元上处理。

张量并行主义。 它涉及将模型的层分片到多个GPU上。对于LLM，张量并行主义至关重要，当模型的单个层对于单个GPU来说太大时。这使得层内的巨大矩阵操作可以被多个GPU并行处理。通过张量模型并行主义，模型的单个层可以在多个设备上进行划分。

3 Motivation and Main Idea

长期上下文LLM的演变中出现了显著的激增，上下文窗口从2K扩展到令人印象深刻的256K仅仅是在一年内。这种进展继续进行，LLM的应用现在要求支持甚至更长的上下文。

上下文长度的增加对LLM服务系统带来了巨大的挑战，特别是在上下文窗口扩展到令人印象深刻的大小时，对于KV Cache的计算和内存要求也在不断增加。

表1显示了这种趋势，显示了KV Cache大小与LLaMA2-13B模型的上下文长度增长的直接对应关系。当前的GPU具有从几十GB到几TB的内存容量，正在受到限制，需要更多的内存空间来容纳KV Cache不断增长的大小。

Challenges to LLM serving on Cloud

在本文中，作者努力有效地利用数据中心中GPU和CPU的大量内存容量，目标是创建一个专门为能够处理非常长上下文长度的LLM服务而设计的有效内存池。然而，开发这样一个系统并非易事。作者已经确定了在云平台上为扩展上下文的大型语言模型提供服务时出现的主要两个挑战。

挑战1：内存需求之间的巨大差异阻碍了模型并行主义。与在整个自动生成过程中不断扩展的KV Cache不同，剩余激活张量的内存需求仍然保持不变，如表1中所述。

注意层与其他层之间的这种差异对于有效实现模型并行主义构成了一个巨大的挑战。为了满足长上下文任务所需的广泛KV Cache，需要更多的GPU。但是，其他层的张量维度不会随着上下文长度的增加而扩展。因此，传统的模型并行主义在跨更多GPU分布时，会导致这些层在这些GPU上进行更细粒度的子划分，如图1所示，从而导致更少的有效资源利用。

一些先前的研究提出将KV Cache划分为更小的块以实现更细粒度的内存管理，以防止内存碎片。虽然这些方法将注意力层的KV Cache与Transformer块分离，但它们仍然依赖于将所有块在本地GPU内存中收集和定位以执行注意力模块的计算。相比之下，作者的设计目标关注于以分布式的方式存储KV Cache并执行注意力模块，这对于在云环境中有效利用可用的丰富资源至关重要。

挑战2：KV Cache大小的动态性导致在云环境中的资源管理低效。自动回归设计的基本性质决定了在生成过程结束之前，最终序列长度仍然是未知的，通常由一个“结束”字符Token。因此，内存需求是完全动态且不可预测的，在规模上波动显著。需求范围可以从几GB到几百GB，并且仍在进一步增加。

这种变化性排除了任何形式的预先资源规划。资源必须动态分配或释放，以应对自动回归过程的实时需求。如果一个上下文所需的内存超过了实例的GPU容量，整个任务必须转移到具有更多GPU的更大实例，这是一个称为实时迁移的过程。实时迁移是资源密集型的，正如作者的实验所示，它可以比标准推理的成本高出25倍。另一种选择是从一开始就为计算实例分配更多的GPU，这可能导致涉及较短上下文的任务造成资源浪费，从而使有效资源分配的挑战更加复杂。

PagedAttention 通过使用类似操作系统的页面细粒度子块来管理KV Cache。然而，这种方法仅限于使用CPU内存进行交换，这种方法在云端上证明是不高效的。由有限CPU内存所施加的限制不仅限制了LLM服务可支持的上下文长度最大值，而且没有利用云中分布的广阔内存资源。相比之下，作者的目标是充分利用数据中心中GPU和CPU的广泛内存能力。作者旨在建立一个精心打造的、为LLM服务设计的有效内存池，以支持处理极其长的上下文长度。

Main Idea

受到以上挑战的驱动，作者提出了一系列关键技术，专门用于解决这些挑战。这些技术共同形成了一个全面和系统的方法，以确保高效的LLM服务，能够处理扩展的上下文长度。

为了应对挑战1，作者提出了一种名为DistAttention的新注意力算法。这种算法将传统的注意力计算分解为更小、更易管理的单元，称为宏注意力（MAs）及其相应的KV Cache（rBlocks）。这种创新方法使得KV Cache的计算与标准Transformer块的标准 Transformer 块分离，从而允许独立地采用模型并行主义策略和注意力层与其他层之间的内存管理。对于非注意层，应用已建立的模型并行主义策略。相比之下，注意层则由适应性管理，动态地在数据中心分配内存和计算资源，以响应KV Cache的变化。

为了克服挑战2，作者提出了DistKV-LLM，这是一种与DistAttention无缝集成的分布式LLM服务引擎。DistKV-LLM旨在提供一种高效的KV Cache管理服务，协调数据中心内GPU和CPU的内存使用。当一个LLM服务实例由于KV Cache的增加而遇到内存短缺时，DistKV-LLM可以主动识别并借用其他实例（具有过剩容量）的可用内存空间，如图2所示。

这种自动机制是通过rManager和gManager这两个主要组件的协作实现的。rManager在每个LLM服务实例内虚拟化所有GPU和CPU内存，处理来自本地和远程实例的内存操作请求。同时，gManager作为全局协调器，维护一个协议，以确保分布式rManager之间的有效、可扩展和一致的资源管理。

此外，DistKV-LLM提出了一个新的算法DGFM，有效地解决了数据中心分布式KV Cache环境中的内存碎片问题。这种共同努力确保了连续和高效的内存利用，从而提高了LLM服务的整体性能和可靠性。

总之，作者的集成方法与DistKV-LLM和DistAttention相结合，为长上下文LLM在云上的服务提出了一个强大和可扩展的解决方案。通过解决与内存管理和计算调度相关的主要问题，确保LLM服务可以在云上高效和自适应地运行。这个创新框架不仅优化了资源利用率，还为实现大规模语言模型部署的未来先进技术铺平了道路。

4 Method

Overview

在接下来的部分中，首先介绍DistAttention，这是一种创新注意力算法，旨在促进分布式KV Cache管理和计算，详细内容请参见第4.2节。基于此提出DistKV-LLM，这是一种专门为集群规模的分布式GPU内存高效KV Cache管理的LLM服务引擎。

作者的方法包括几个关键组件：首先，在第4.3节中介绍了rManager，这是一种为每个LLM服务实例虚拟化GPU和CPU内存的软件层。它提供了一个抽象层，称为rBlocks，以提高内存管理效率。其次，在第4.4节中描述了一个由gManager启用的全面协议，这是一个全局管理系统。它确保了数据中心内分布式GPU上的rBlocks的有效、安全和一致管理。在第4.5节中进一步提出了一种专门设计来聚合碎片化内存块的创新算法，从而显著提高分布式KV Cache系统内的数据局部性。最后，在第4.6节中提出了一系列优化策略，旨在最小化与KV Cache分布式存储相关的巨大通信开销，通过有效地重叠计算和通信任务。关于这个设计的更多细节将在以下部分中讨论。

DistAttention

为了应对内存管理的复杂性，作者开发了一种新的注意力算法DistAttention。这种算法有效地将KV Cache划分为较小、更易管理的子块，从而促进数据中心内的分布式内存管理。这种方法的关键在于将DistAttention划分为多个微观注意力（MAs），每个MA都包含KV CacheToken的一个子序列。这种设计设计的一个独特方面在于，它可以通过分别计算MAs来计算注意力结果。当所有微观注意力（MAs）在其各自的Token子块上完成计算后，通过汇总过程可以得到最终注意力结果。这个汇总过程涉及缩放和降低每个MA的输出。这个汇总过程的详细方法和精确公式将在以下方程中详细说明：

MA_{ij}=exp(Q_{i}K_{j}^{T}-max(Q_{i}K_{j}^{T})) \tag{3}

Attention(Q,K,V)=Reduce(Scale([MA_{ij}]_{j=1}^{B_{U}})) \tag{4}

其中Reduce的计算如下：

\begin{split}& Reduce(Scale([MA_{iI}]_{j=1}^{B_{kv}}))\\ &=Reduce([exp(max(Q_{i}K_{j}^{T})-max_{i})MA_{ij}]_{j=1}^{B_{kv}})\\ &=\sum_{j=1}^{B_{kv}}(exp(max(Q_{i}K_{j}^{T})-max_{i})MA_{ij}/sum_ {i})\\ & max_{i}=max(max(Q_{i}K_{1}^{T}),...,max(Q_{i}K_{B}^{T}))\\ & sum_{i}=\sum(exp(Q_{i}K_{j}^{T}-max_{i})\end{split} \tag{5}

这种方法不仅将单个MA执行的计算进行整合，而且还有效地将它们合成一个连贯的最终输出，展示了作者实现在注意力算法中分布式处理框架的有效性。

The rBlock and rManager

在实现DistAttention后，LLM的关键值（KV）缓存被分割成更小的单元，称为rBlock。每个rBlock包含一组对应于固定数量的关键值（Key-Value）Token的向量，以及一些关键的元数据。这些元数据提供了关于子块的 critical 信息：rBlock ID和实例ID指示该rBlock中的KV Cache是否属于本地实例还是远程实例；设备和物理ID标签了rBlock的物理位置，它可以位于CPU侧或多个GPU之一。

每个LLM服务实例都配备了一个专门的rBlock Manager，称为rManager。rManager负责监控本地设备中的所有rBlock。它通过将全局GPU内存划分为固定大小的物理rBlock来实现对全局内存空间的虚拟化。每个物理rBlock都设计为一个逻辑rBlock的容纳单元。rManager维护了一个详细的映射表，将这些逻辑rBlock映射到全局内存中相应的物理rBlock地址，如图3所示。

rManager提供了一个统一的API接口，既可以服务本地实例，也可以服务远程实例。这些操作包括为新生成的KV Cache分配物理rBlock，并在不再需要时释放它们。当收到内存分配请求时，rManager会 Query rBlock表以确定第一个可用的物理rBlock空间。

在场景中，如果可用空间不足，rManager将启动从其他实例中借空间的程序。更多细节将在第4.4节中详细说明。值得注意的是，如果分配请求来自远程实例，rManager被编程为自动返回一个错误的响应，表示无法直接远程分配。作者对可以分配给远程实例的rBlock数量进行了限制，该数量是通过实验确定并配置为超参数。

The gManager and Contract Protocol

作者系统的关键组件，称为gManager，作为中央管理者，负责维护所有实例的全球内存信息。每个实例定期向gManager发送心跳信号，以传达其剩余可用内存空间的更新。利用这些数据，gManager构建了一个详细的表格，称为全球债务账簿，如图4所示。

当某个实例的内存空间不足以支持rBlocks时，对应的rManager会向相邻的实例借用GPU或CPU内存空间。在此之前，rManager作为债务人，需要向gManager发起 Query ，告知需要借用的内存空间大小。当收到这个 Query 时，gManager会查阅全球债务账簿，并回应潜在的债权人ID，这些ID代表目前拥有剩余内存空间的实例。选择过程遵循地方性和可用性原则，即选择债权人实例时，基于最低的相对通信成本和最高的可用内存空间。gManager提出了三个潜在的债权人ID作为建议。

随后，债务人实例会依次向这些债权人实例发出请求，直到成功确认分配。在所有候选人返回否定回答的情况下，债务人实例会回退到gManager寻求其他建议。这种动态和响应式系统确保了数据中心内内存空间的分配和管理高效有效。

接下来，将描述合同协议的关键组件和设计考虑因素。

全球债务账簿：由gManager管理的全球债务账簿是一个关键的表格，记录了网络中可用的内存和实例之间的债务。每个条目都代表一个LLM服务实例，详细列出了实例ID和可用的内存空间。每个条目中的后续子字段标明了债务人的实例ID以及它们各自从债权人那里借来的rBlock数量。

竞争候选人：在多个债务人实例同时向rManager发送请求的场景中，系统必须高效地处理这些竞争需求。全球债务账簿在此处发挥着重要作用，使gManager能够均匀地将请求分配给实例，从而防止任何单个实例过载。另一方面，rManager采用先到先服务的策略，从远程实例为rBlocks分配物理空间。如果由于空间限制，rManager无法为远程rBlocks分配足够的物理rBlock，它会向债务人实例发送一个虚假的响应。这种响应也会促使gManager更新当前的资源可用性记录，从而在有可用资源之前暂停将新请求转发。这种方法确保了内存资源的平衡有序分配，减少了系统的潜在瓶颈。

一致性：采用了一种松散的一致性策略，即gManager和rManagers之间。在这种方法下，gManager不需要严格跟踪所有实例的内存分配或释放操作。相反，它通过定期自动发送的心跳来收集这些信息。因此，gManager维护的是数据中心内的总体空间使用情况，而不是详细、实时的数据。

当rManager的债务人请求借用空间时，gManager只提供潜在债权人候选人的建议。债务人必须与这些建议的债权人进行谈判以确定内存分配。涉及多个同时向同一rManager发出请求的情况，使用之前讨论的竞争候选人策略进行管理。这种松散耦合的一致性框架不仅简化了操作，还最大限度地减少了过度交易开销，从而降低了处理延迟并提高了整体系统性能。

可扩展性：为了满足不同的吞吐量需求，gManager被设计为通过并发处理 Query 请求来提高可扩展性。为了加速确定具有剩余内存空间的实例的过程，gManager定期启动排序操作。这个操作根据实例的剩余可用内存空间对实例进行排序，使得 Query 请求可以有效地绕过具有最小内存资源的实例。这种方法确保gManager在其最优容量下运行，在扩展以满足网络动态需求的同时，保持系统效率和响应性。

Fragmented Memory Management

由于上下文长度的动态性和批处理，内存管理面临着一个关键挑战，即碎片化。系统中的每个实例都既是内存的债权人又是债务人，根据需要从其他实例借贷或释放内存。例如，处理长上下文请求的实例可能会持续增长，需要从远程实例借空间。

相反，处理短命请求的实例会尽早释放内存空间，然后可以将空间借给其他实例或分配给新请求。这种动态性导致了一个重要问题：数据局部性的恶化。由于实例经常访问远程内存位置的数据，系统会受到显著的性能惩罚，例如增加延迟和降低吞吐量。

作者提出了一种基于债务图的碎片化内存管理算法，称为DGFM，旨在通过有策略地召回已借出的内存空间并将其交换为本地数据存储来应对这个问题。该问题的一个关键挑战是数据中心中运行的大量LLM服务实例，这些实例经常涉及复杂的债务关系。

为了有效地管理这种复杂性，作者将这个问题概念化为在有向图中搜索圆。首先，作者构建一个反映这些债务关系的有向图镜像，其中每个节点表示一个实例，每条有向边表示一个实例对另一个实例的债务。作者的算法然后是迭代应用的。在每次迭代中，算法随机选择一个节点并遍历图以识别一个有向圆。发现这样的一个圆是至关重要的，因为它表明涉及的节点（或实例）可以相互解决他们的债务。这种解决方法是通过有策略地召回和交换各自的内存块来实现的，从而提高整体系统效率和内存利用率。这个算法的详细信息如图5和算法1所示。

该有向图是从全局债务账簿中派生出来的，DGFM算法由gManager执行。当识别到一个有向环时，gManager会向对应实例的rManager发出请求，冻结它们以防止修改或取消。作者为每个节点设置了一个经验阈值，以确定需要交换的最小内存块数量（rBlocks），以防止在过于小的内存块上进行低效的召回和交换操作。这个过程显著减少了远程内存访问的需求，从而提高了数据局部性，最终导致系统性能的明显改善。

Communication Optimization

分布式KV Cache存储面临另一个挑战：在LLM服务中执行长上下文任务时，预填充和自动回归阶段可能生成大量KV Cache，导致rManager向远程空间借用。作者已经针对这两种情况进行了特定的优化，如图6所示。

在预填充阶段，KV Cache的内存需求可以精确预测，基于提示的长度。这种预见性使得可以预先计划rBlocks的分配，根据它们的存储位置，可以Token为本地或远程。当执行Transformer块的注意力层时，作者将注意力计算与远程rBlocks的传输重叠。

在自动回归阶段，rBlocks的分配是动态的。简单地将所有rBlocks用于本地计算会导致过多的网络流量。此外，由于注意力模块的计算本质上是向量-矩阵乘法，这是一个明显的内存密集型任务，将所有计算本地化会严重降低系统性能。DistAttention的创新使作者能够将 Query 向量重定向到包含远程rBlocks的实例中，以便在那里进行宏注意力计算，然后再将结果发送回进行集成。这种方法将数据传输量减少了

倍，其中

表示KV Cache中的Token数量。这种方法的局限性在于，它可能会与远程rBlocks的主实例争夺计算资源。

为了减轻这种影响，在每个rManager中建立了一个阈值，根据本地SLA指南，根据借用计算资源的数量进行裁决，以确保系统计算资源的平衡和谨慎使用。

5 Implementation Details

DistAttention包含两种类型的运算符，分别是DistAttn和ScaleReduce，这些运算符的开发使用了大约5000行的C++/CUDA代码。DistAttn运算符是专门用于分布式注意力计算的，通过ScaleReduce运算符进行结果汇总，从而得到最终的输出结果。

为了适应输入上下文长度的范围，DistAttn采用了一种自适应核选择过程，该过程基于输入的尺寸。上下文长度被分为三个组：正常范围（0-8k）、长范围（8k-32k）和超长范围（>32k），作者对这些组分别开发并精心优化了三个不同的核模板。此外，作者还设计并实现了一种启发式方法，用于针对特定张量形状的自适应CUDA核的微调。

另一方面，DistKV-LLM将Ray框架调整到分布式KV Cache管理和调度系统，这个系统使用了大约12000行Python代码进行开发。为了有效地实现请求和网络数据传输，作者自定义了包编码，并使用套接字而不是基于RPC的框架来传输数据包。为了最大限度地利用RDMA的高带宽优势，作者使用了NCCL进行跨实例GPU张量通信，确保分布式实例之间的数据交换高效。

6 Evaluation

在这一部分将介绍作者工作的评估结果。

环境。作者在一个拥有4个节点和32个GPU的集群上部署了DistKV-LLM。每个节点有8个NVIDIA A100（80GB）GPU。

模型。作者的框架现在可以支持大多数流行的LLM，如GPT，LLaMA，BLOOM等。由于大多数LLM模型具有相似的Transformer块，作者选择一个代表性的模型LLaMA2进行评估。LLaMA2家族包含三个不同的模型大小：7B，13B和70B。它们使用两种流行的注意力架构；7B和13B模型使用多头自注意力（MHA），而70B模型使用分组 Query 注意力（GQA）。

基准。作者选择vLLM，即最新的LLM服务引擎，作为主要的基准。此外，大多数先前的LLM服务系统使用张量并行主义。为了验证流水线并行主义与连续批处理，在vLLM框架中实现了一种与Alpa相似的设计，作为基准之一。

Context Length Benchmark

在不同的上下文长度下评估和比较DistKV-LLM和基准的性能。作者评估了三个具有不同上下文范围模型。对于LLaMA2-7B，在2个GPU上评估了1-200k的任务，在16个GPU上评估了1-1200k的任务，在32个GPU上评估了1-1900k的任务。对于LLaMA2-13B模型，在4个GPU上评估了1-280k的任务，在8个GPU上评估了1-480k的任务。对于LLaMA2-70B模型，在8个GPU上评估了1-450k的任务。

为了验证DistKV-LLM的性能，作者与两个vLLM基准版本进行了比较。vLLM-v1在单个实例中包含与DistKV-LLM相同数量的GPU。

图7显示了vLLM-v1，vLLM-v2和DistKV-LLM在不同上下文长度下的吞吐量。值得注意的是，DistKV-LLM（蓝色）不仅实现了与vLLM-v1（绿色）相当的吞吐量，而且支持更长的上下文长度，如图7所示，大约为2x-19x。这种改进归因于DistKV-LLM在所有实例之间有效地协调内存使用的能力，而vLLM-v1仅限于实例的私有内存。

vLLM-v2预先分配了更多的GPU，以便支持与DistKV-LLM相当的上下文长度。通过与vLLM-v2（红色）的比较，作者证明了DistKV-LLM在保持相似的扩展上下文长度的同时，实现了显著更高的吞吐量。

如图7所示，DistKV-LLM比vLLM-v2实现了1.4x-5.3x的更高吞吐量。这是因为DistKV-LLM可以保持高效的模型并行主义策略，而vLLM-v2将模型在更多GPU上划分成较小的片段，导致硬件效率较低。

End-to-end Serving Performance

作者采用了6.1节中的实验设置，运行相应的上下文范围数据集来评估DistKV-LLM的端到端性能。实验结果如图8所示。当曲线急剧上升时，表示吞吐量已达到系统限制，请求开始积压，延迟迅速增加。

在包含1%长请求的数据集中，DistKV-LLM相对于基准实现了大约1.4x到2.4x的改进。这是因为将模型分成更小的片段导致GPU利用率降低，从而显著降低了线性计算的效率。在包含10%长请求的数据集中，DistKV-LLM相对于基准实现了大约1.1x到1.4x的性能改进。在一个包含30%长请求的数据集中，DistKV-LLM相对于基准实现了大约1.03x到1.3x的性能改进。

随着数据集中长请求所占比例的增加，DistKV-LLM提供的性能提升逐渐减小。这是因为当模型处理长上下文请求时，线性计算与注意力计算的比例较低。DistKV-LLM在线性计算部分获得的性能提升成为整体计算工作量较小的较小部分，而注意力计算的性能没有显著优于基准。

Live Migration

另一种解决上下文长度变化的方法是实时迁移，它可以在需要时将任务迁移到具有更多GPU的更强大的实例上。在本实验中，将比较DistKV-LLM和实时迁移在LLaMA2-7B模型上。对于新实例，LLM模型通过Amazon简单存储服务（Amazon S3）下载，并以SafeTensor格式由vLLM加载。

最初，使用一个A100 GPU来部署服务，它可以处理最大长度为108k的请求。当上下文长度超过108k时，应使用额外的A100 GPU进行扩展。结果如图9所示。

水平轴表示提示的长度，垂直轴表示生成对应输出长度时的延迟。由于实时迁移导致的开销是DistKV-LLM中通信开销的45倍。当上下文长度为105k提示和5k输出时，会触发实时迁移。在这种场景中，vLLM的延迟显著增加，而DistKV-LLM仅经历微小的干扰。当生成长度为5k，10k，20k和30k的Token时，DistKV-LLM的完成时间分别是vLLM的3.5倍，2.2倍，1.6倍和1.4倍更快。这是因为迁移整个服务会产生大量的开销，包括通过高速下载链接远程下载模型（尽管使用了高速下载链接）和推理引擎加载模型。相比之下，DistKV-LLM只需要与扩展设备建立连接，无需下载和加载模型。

Optimizing Memory Allocation

上下文长度和批处理的动态性导致了数据局部性的恶化。DGFM算法旨在通过调用借出的内存空间来优化内存分配。在本实验中使用具有四个LaMA2-13B tp2实例的DistKV-LLM服务，能够处理请求长度从1到480k的请求。作者比较了启用DGFM和未启用的服务的吞吐量性能，结果如图10所示。

在初始阶段，启用DGFM和未启用的服务的性能相似。随着时间的推移，数据局部性问题开始出现。具有DGFM的服务通过定期清除债务圈和优化分布式环境中的整体数据局部性，保持了更高的整体吞吐量。相比之下，未启用的服务出现了越来越严重的问题，导致整体吞吐量持续下降。

Ablation Study

DistAttention Breakdown。DistAttention引入了在实例之间分布式存储和计算注意力的能力，这也会带来一些开销。这些开销主要是由于单个Token的 Query 、键和值张量的传输成本，以及一些其他开销，如张量切片和拼接。

在8个A100 GPU的实例上部署了两个实例，通过DistKV-LIM共享存储和计算资源。对DistAttention的运行时间进行了 breakdown 分析，并将其与通过张量并行主义将注意力分为八部分的注意力进行了比较。

结果如图11所示。与TP8 Attention相比，额外的开销引入了5%-10%的延迟。这额外的延迟几乎恒定，这意味着随着上下文长度的增加，这个开销比例越来越小。

比较远程计算和本地计算。有两种策略来计算远程分配的rBlocks，

本地计算：通过高速互联网络将远程实例的KV Cache带回来，在本地执行完整的注意力计算；
远程计算：将 Query 向量传输到远程实例所在的rBlocks位置，通过DistAttention实现分布式注意力计算，然后将结果向量取回。这两种方法的结果如图12所示。

本地计算的延迟明显高于远程计算，这是因为本地计算需要通过网络将大量远程KV Cache传输回本地实例，受到网络带宽的限制，极大地增加了整体延迟。在DistKV-LLM中，作者使用这个实验结果在rManger中指导第4.6节中讨论的通信优化。

7 Related Works

现有的LLM服务系统。最近提出了许多LLM服务系统。ORCA引入了一个迭代 Level 的调度策略，极大地提高了批处理的计算和内存利用率。为了解决由于碎片化和重复复制导致的内存浪费问题，vLLM开发了一个Paged KV缓存和Paged Attention机制。DeepSpeed-FastGen提出了一种新颖的提示和生成组合策略，称为Dynamic SplitFuse，旨在进一步增强连续批处理和系统吞吐量。AlpaServe探索了在突发请求率场景下通过模型并行进行统计复用的机会。

FasterTransformer和DeepSpeed Inference实现了针对Transformer模型进行 pioneer 和 extensive Kernel Level 性能优化的 pioneered 和 extensive。TGI, TensorRT-LLM和lmdeploy在基于FasterTransformer的基础上，适应了连续批处理和分页注意力等特性。尽管这些新颖系统解决了许多问题并取得了出色的结果，但动态问题以及需要支持异常长的上下文长度仍然是一个未解决的挑战。

与环注意力（Ring Attention）的比较。环注意力（Ring Attention）是一种将长序列分布在多个设备上的方法，旨在使块级注意力与KV块的通信完全重叠。这种方法对于长序列的训练和推理预填充阶段非常高效。然而，在推理的解码阶段，设备之间KV块的传输不能通过计算来隐藏，导致大量的开销。图12显示了KV块传输的开销显著高于DistAttention的通信开销。

解决长上下文推理管理超大数据值（KV）缓存的方法的另一类涉及稀疏KV缓存，例如滑动窗口注意力（Sliding Window Attention）。这种技术只需要维护窗口大小的KV缓存。H2O和StreamingLLM也保留了一个固定的窗口大小KV缓存，但它们通过采用KV缓存删除算法和集成注意力 sink来缓解由于上下文信息丢弃而导致的精度损失。然而，由于这些方法会丢弃一些上下文信息，因此它们在一定程度上会损害大型语言模型（LLMs）的有效性。

8 Conclusion

LLM推理计算的动态、自回归性质对云端的LLM服务带来了显著的挑战，尤其是在处理长上下文序列的任务时。为了解决这些挑战，作者提出了DistAttention，这是一种创新的分布式注意力算法，可以有效地将KV缓存分割成可管理的单元进行分布式处理。

此外，作者还提出了一种分布式LLM服务引擎DistKV-LLM，它在KV缓存管理方面表现出色，可以优化数据中心内内存使用。这种结合DistAttention和DistKV-LLM的方法有效地确保了在处理长上下文任务时，LLM云服务的顺畅和高效。在32个NVIDIA A100 GPU和18个数据集的综合评估中，作者的系统显示了比现有解决方案1.03-2.4倍性能改进，并支持了2-19倍更长的上下文长度，证明了在管理各种上下文生成任务方面的有效性。

参考

[1]. Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2024-04-11，如有侵权请联系 cloudcommunity@tencent.com 删除

阿里巴巴

本文分享自 AIGC 先锋科技微信公众号，前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！