• 文章标题:利用模型-注意力解耦实现高效的异构大语言模型解码
  • 作者/机构:Shaoyuan Chen1, Wencong Xiao2, Yutong Lin1, Mingxing Zhang1, Yingdi Shan1, Jinlei Jiang1, Kang Chen1, Yongwei Wu1
  • 机构:1清华大学 2字节跳动

A1 主要贡献

本文针对基于Transformer的大语言模型(LLM)在实际服务中,由于昂贵的计算优化型加速器使用效率低下而带来的挑战,提出了一种创新的解决方案。现有研究虽然提出了将LLM推理拆分为不同阶段的解耦服务架构,但解码阶段的效率仍然很低,其根本原因是Transformer模型中不同算子的资源需求不同。具体来说,注意力(attention)算子是内存密集型的,其内存访问模式与现代加速器的优势相冲突,尤其是在处理长上下文请求时。

为了提升LLM解码效率,本文提出了模型-注意力解耦(model-attention disaggregation)的核心理念。
- 研究目标与核心理念:该方法利用一组廉价的、内存优化的设备来专门处理注意力算子,同时继续使用高端加速器处理模型的其他部分。这种异构设置确保每个组件都能针对其特定的工作负载进行优化,从而最大化整体性能和成本效益。
- 可行性分析:通过全面的分析和实验,本文证实了在多个设备上拆分注意力计算的可行性,并证明异构设备之间所需的通信带宽在现有网络技术的可控范围内。
- 系统实现:为了验证理论,本文开发并部署了Lamina,一个在分布式异构集群中集成模型-注意力解耦的LLM推理系统。Lamina通过一系列技术优化来应对解耦带来的挑战:
1. 全主机旁路网络栈(fully hostbypassed network stack):利用PCIe P2P功能,使GPU能直接与网卡(NIC)通信,消除了主机CPU的同步和参与,从而降低网络开销。
2. 自动化模型转换器(automated model converter):该转换器将模型的计算图分割成与注意力算子交错的切片,并重新排序算子,协调计算和通信流水线,以实现通信和计算的有效重叠。
3. 交错流水线(staggered pipelining):通过并发运行多个批次并优化工作流,确保计算设备和内存设备能同时工作,从而最小化资源浪费并最大化系统性能。
- 实验成果:实验结果表明,在硬件成本相近的情况下,Lamina的预估吞吐量比现有解决方案高出16.1%至90.1%。

A3 背景知识:LLM解码中GPU的低利用率

本节对当前同构硬件下LLM解码实现的挑战和限制进行了详细的性能分析,以LLaMA3-70B模型为例。分析中使用的具体符号见表2。

表2:性能分析中使用的符号。同时给出了LLaMA3-70B模型对应的值。

2.1 预备知识

Transformer架构基础。现代LLM主要依赖于Transformer架构【索引49,Attention is all you need,2017,NIPS】。每个输入词元首先被映射到一个维度为d的词嵌入。这些嵌入随后通过一系列Transformer模块。最终的输出嵌入与一个采样矩阵相乘以生成下一个词元的预测概率。

Transformer模块内的计算。在每个Transformer模块中,输入嵌入被投影成三个独立的向量:查询(qi)、键(ki)和值(vi),它们的维度与隐藏状态d相同。这些向量通过一个注意力算子来计算注意力分数。然后,注意力分数通过一个矩阵Wout进行加权,以产生注意力层的输出嵌入yi。

前馈网络(FFN)。输出yi接着通过一个前馈网络,该网络将其扩展到一个中间向量空间,然后通过另一次矩阵乘法将其缩减回来:

两种核心计算操作。尽管Transformer模块涉及多种变换,但计算密集型操作实际上只有两种:注意力算子(在公式中用⋆表示)和其他矩阵投影步骤。因此,本节将基于屋顶线模型【索引50,Roofline: An insightful visual performance model for multicore architectures,2009,Commun. ACM】和实验测量,对这两种算子进行定量分析,以突显它们在解码阶段的不同特性,解释为何当前同构硬件的LLM解码实现常导致GPU利用率不足,从而引出对异构架构的需求。

2.2 硬件利用率不足

2.2.1 非注意力算子的利用率不足

通过连续批处理提升利用率。为了提高LLM解码中的GPU利用率,连续批处理技术被广泛采用【索引16,Lazy batching: An sla-aware batching system for cloud machine learning inference,2021,HPCA;索引20,Low latency rnn inference with cellular batching,2018,EuroSys;索引46,E-batch: Energy-efficient and high-throughput rnn batching,2022,ACM Trans. Archit. Code Optim.】。通过并发处理多个输入,GPU内存中的模型参数得以重用,使得工作负载变得更加计算密集。对于一个包含B个请求的批次,非注意力算子大约需要2NB次浮点运算。此外,这些算子需要加载模型参数eN,并从GPU内存中读/写总共2eBd的输入和输出数据。由此产生的算术强度为 $2NBe(N+2Bd)$ ,它会随着批次大小的增加而迅速增长。

算术强度与硬件利用率的权衡。图2展示了在NVIDIA H100 GPU上测量的LLaMA3-70B非注意力算子的延迟和内存吞吐利用率(MFU),以及基于屋顶线模型的预测值。对于小批量(小于100),工作负载是带宽受限的,延迟主要由从GPU内存访问模型参数引起。在这种情况下,MFU低于20%,表明计算资源存在严重的未充分利用。随着批次大小增加,工作负载转为计算受限,延迟也随之增加。为了优化GPU资源利用率,通常倾向于使用更大的批次。然而,这常常受到有限的VRAM容量的制约,因为VRAM无法容纳所需的大量KV缓存,这一限制将在后面详细讨论。


图2:在LLaMA3-70B模型单次解码迭代中,非注意力算子的实测耗时和MFU。展示了不同张量并行度下的结果。虚线表示使用屋顶线模型的预测值。

2.2.2 注意力算子的利用率不足

注意力算子的内存密集特性。与权重矩阵投影算子不同,注意力算子在处理一批请求时,仍然执行的是批处理矩阵向量乘法(BGEMV),其中每个查询都访问并处理其自身的KV缓存。因此,注意力算子的算术强度保持不变,与批次大小无关。这种特性使得注意力操作成为内存受限的,并且增加批次大小并不能提高资源利用率。较新的模型采用了分组查询注意力(GQA),它将qi分成G个独立的查询组,并将ki和vi的大小减小G倍。每个查询与相同的ki和vi进行注意力计算,然后将输出简单地拼接起来。通过GQA,注意力算子的算术强度增加了G倍,但与其他算子相比仍然相当低。

内存容量成为瓶颈。如图3所示,即使对于像20这样的小批量,注意力算子的带宽利用率仍然保持在70%以上。即使在像H20这样的内存专用加速器上也是如此,H20的TFLOPs性能仅为H100的15%。然而,注意力操作可实现的批次大小受到GPU内存容量的限制,特别是由于长上下文长度下KV缓存的高内存需求。例如,在上下文长度为8192时,一块H100的全部内存只能容纳大约30个请求的KV缓存,而由于模型权重占用了部分内存,实际数量会更低。因此,注意力操作的有限批次大小成为一个关键瓶颈,阻碍了在解码阶段对非注意力操作计算资源的有效利用。


图3:在LLaMA3-70B模型单次解码迭代中,注意力算子的实测耗时和模型带宽利用率(MBU)。展示了不同序列长度和硬件配置下的结果。

A2 方法细节

3.1 概述

同构方案的资源利用问题。当前的LLM服务系统在解码阶段通常对注意力和非注意力算子使用相同的硬件。然而,我们的分析揭示,这种同构方法会导致两种算子的资源利用率都 suboptimal,原因如下:
- 注意力算子的算术强度低,因为从KV缓存中检索的每个值只参与有限次数的计算。鉴于现代高性能加速器中内存带宽和计算能力之间的差距(偏好高算术强度的负载以实现高效资源利用),这些算子往往会低估先进GPU的计算资源。
- 对于非注意力算子,虽然增加批次大小可能提高硬件利用率,但这也会导致KV缓存相应增加,可能超出可用内存容量。因此,为防止内存溢出,批次大小通常保持较小,这同样因为算术强度低而导致硬件利用率低下。

模型-注意力解耦架构。为了解决同构解码方案的上述限制,我们提出了模型-注意力解耦架构。该架构使用内存专用加速器来存储KV缓存并计算注意力算子;非注意力算子则仍在原始加速器上执行。一个模型-注意力解耦系统可以使用多种设备的多个实例来提供不同程度的并行性(DOPs)。如果我们用a个GPU处理非注意力算子,b个内存优化型GPU处理注意力算子,我们将此DOP表示为(a, b)。

解耦架构的优势。通过利用更便宜的内存优化设备,我们可以支持更大的批次大小,因为扩展的内存容量可以存储更多的KV缓存,从而提高了非注意力算子的算术强度和硬件利用率。此外,由于注意力计算被移至内存优化设备,我们避免了浪费高端GPU宝贵的计算资源。

通信开销分析。实现注意力卸载的一个潜在障碍是,模型每一层都需要在异构加速器之间进行数据传输,这可能遭遇“通信墙”问题并增加端到端的解码延迟。我们进行了定量分析,以确定此类传输所需的互连带宽。假设我们以批次大小B运行一次迭代,并且可以为网络开销承受α倍的额外延迟,那么所需的最小互连带宽可以计算为:


其中,$M_{TIME}(B)$ 和 $A_{TIME}(B,l)$ 分别是在批次大小为B和序列长度为l时,非注意力和注意力算子的运行时间,这些都可以通过实验测量得到。当α = 0.2时,针对不同批次大小计算出的所需最小带宽估算值如图4所示。

网络带宽需求是可控的。从数据中可以明显看出,即使在处理高达300的批次时,所需的互连带宽也不超过30 GB/s。这种带宽需求可以通过像400Gbps以太网这样的网络技术轻松满足。事实上,当代数据中心已经满足了这一要求,其中每个GPU通常都配备一个专属的400Gbps NIC,为LLM训练提供充足的网络带宽。


图4:使用H100和H20进行注意力卸载解码LLaMA3-70B所需的网络带宽,网络开销导致的延迟减慢不超过20%。

内存设备的网络配置。对于内存设备,同样需要相同的互连带宽来与计算设备通信。由于我们采用了一组更经济但性能较弱的内存设备来协同计算注意力,因此每个独立设备所需的通信带宽要小得多。因此,我们可以选择为每个设备配备一个性能较弱的NIC,或者安装一个共享的400Gbps NIC来服务多个内存设备。

3.2 实践挑战

模型-注意力解耦的挑战。尽管模型-注意力解耦在提高LLM解码效率方面展现了潜力,但它也带来了一系列严峻的实践挑战,具体如下:
- 频繁的网络通信:将注意力算子从计算优化设备分离到内存优化设备,引入了每个模型层内的跨机器数据通信。尽管现有数据中心的互连带宽足以支持注意力卸载,但我们发现网络延迟可能仍是高效LLM解码的一个问题。通过注意力卸载,我们在不同节点上的GPU之间进行逐层数据传输,这可能导致每秒高达数千次的往返。这些频繁的网络传输由于累积的网络延迟,可能会显著增加解码时间。因此,需要一个经过改造的、延迟优化的、GPU感知的网络栈,以实现模型-注意力解耦的最佳性能。
- 软件工程挑战:通过模型-注意力解耦,我们将Transformer模块的中间操作——注意力算子的执行移至其他设备。这需要对现有的LLM代码库进行复杂且破坏性的修改。具体来说,我们必须将模型分解成与Transformer-based LLM的模块化结构不符的独立切片。这个过程不仅劳动密集且容易出错,还显著增加了维护复杂性。因此,非常需要自动化工具来帮助切分模型并执行相关优化。
- 执行重叠困难:在异构解耦系统中,各种设备如计算优化型GPU、内存优化型GPU和NIC可以同时被利用。因此,如果占用不同设备的操作能够重叠执行,我们可能会实现显著的执行时间缩减。然而,在当前LLM的Transformer架构中,注意力和模型算子以顺序方式紧密交错,一个算子的输出通过网络传输成为另一个算子的输入。因此,依赖于不同硬件资源的操作无法有效在时间上重叠,导致相当大的资源利用不足。因此,需要仔细编排各种设备上的操作,并高效设计任务流水线,以促进执行重叠并提高资源利用率。

4. Lamina系统设计

Lamina系统概述。我们构建了Lamina,一个分布式的异构LLM解码系统,它实现了模型-注意力解耦并解决了相关挑战。Lamina采用两种加速设备:内存设备用于存储KV缓存和计算注意力算子,计算设备用于存储模型参数和计算模型的其他部分。这两种设备通过高速数据中心网络(DCN),如Infiniband或以太网互连。

4.1 全主机旁路网络栈

传统GPU网络通信的开销。跨节点GPU之间的通信通常利用RDMA技术,这是一个需要CPU、GPU和NIC等多个系统代理协调的复杂过程。为了减少GPU感知网络的开销,GPUDirect RDMA (GDR)【索引3,GPUDirect RDMA,NVIDIA官方文档】被开发出来,允许RDMA-capable NIC (RNIC)直接访问GPU内存,从而消除了主机内存作为中间缓冲区的需要,提高了网络延迟和带宽。然而,控制路径仍然需要CPU参与,并包括几个步骤,所有这些步骤都位于关键路径上并导致网络延迟。具体来说,使用GPUDirect RDMA传输数据的步骤如下:
1. 本地CPU等待所有先前的GPU内核完成,确保要传输的数据已准备就绪。
2. 本地CPU向RNIC提交一个发送工作请求(WR)。
3. 本地RNIC处理发送WR,从GPU内存中获取数据并通过物理网络链路传输。
4. 远程RNIC接收数据并将其写入GPU内存。
5. 远程CPU等待RDMA接收操作完成。
6. 远程CPU启动后续的GPU内核。
根据我们的实验结果,步骤1到5可能产生60-70 µs的延迟。此外,由于必须在接收数据准备好后才启动内核,GPU内核启动开销(可能高达20 µs)也增加了端到端延迟。所有这些额外的延迟对模型-注意力解耦构成了重大开销,因为它依赖于频繁的网络通信。

FHBN recv实现。为了减少网络开销,我们开发了一个全主机旁路网络(FHBN)栈,它在GPU感知网络的控制和数据路径中完全消除了主机CPU的参与。FHBN的recv函数采用了设备端轮询技术来等待接收操作的完成。具体来说,我们在接收方的GPU内存中分配一个seqno变量。发送方在每次发送操作后通过RDMA write递增远程的seqno。数据发送和seqno递增操作被批量放入一个WR中,因此不会增加端到端延迟。当接收方GPU准备好接收和处理传入数据时,它通过一个专门的GPU内核主动轮询seqno的值。这种方法不仅消除了recv过程中CPU的参与,还允许异步地将轮询内核和后续的计算内核启动到GPU流中,因此,GPU内核启动开销也从关键路径中移除。

FHBN send实现。FHBN send的实现更为复杂,如图5所示,它要求GPU直接向RNIC提交RDMA命令。当CPU向RNIC提交新的RDMA WR时,它首先将WR入队到主机内存中的工作队列(WQ),然后通过敲响门铃(DB)——RNIC的用户访问区域(UAR)中的一个特殊寄存器——来通知RNIC有待处理的工作。UAR是RNIC的MMIO区域的一部分,并映射到非特权应用程序的地址空间,以允许内核旁路的RDMA操作。以上所有步骤都在RDMA用户空间库(libibverbs)中实现。为了在GPU上实现直接的RDMA命令提交,我们必须允许GPU通过PCIe P2P直接访问UAR。具体来说,我们使用cudaHostRegisterIoMemory API将UAR映射到GPU的地址空间。然后,我们在CUDA设备代码中重新实现了RDMA命令提交逻辑。为了进一步降低延迟,我们利用了BlueFlame机制,这是Mellanox RNIC【索引4,Mellanox adapters programmer’s reference manual】提供的一个硬件特性。这种方法允许WR通过对UAR的MMIO写入直接提交给RNIC,从而消除了RNIC通过昂贵的PCIe DMA读取从主机内存获取WR的需要。需要注意的是,WR仍应入队到WQ中,因为硬件可能偶尔会错过BlueFlame WR并回退到常规工作流程,特别是在重负载下。


图5:使用FHBN send和传统GPU感知send实现提交WR的示意图。

4.2 自动化模型转换器

4.2.1 模型切分

切分模型的挑战。在注意力卸载架构中,LLM的不同算子可能在不同的硬件上执行;因此,我们需要将模型分割成切片,这是通过在注意力算子处进行切割来实现的。这通常涉及到对现有代码库进行重大修改,主要是因为期望的切割点与LLM固有的模块化结构不一致。这种不一致性使分区过程复杂化,并增加了异构系统内出现错误和不一致性的风险。

自动化模型切分器。为了方便模型分区,我们开发了一个自动化的模型切分器,能够将LLM转换为可独立调用的切片,如图6所示。给定LLM的源代码,切分器使用符号执行来生成一个加权计算图。每条边的权重表示在算子之间传递的数据大小,这是从模型的形状规范中推导出来的。

最小加权割的应用。由于残差连接和其他复杂的模型结构的存在,直接移除注意力算子并不总能得到一个断开的计算图。因此,我们计算从注意力算子的输入到输出的剩余图的最小加权割。这个最小割中的边代表了在切片调用之间必须保存的上下文,它们将从计算图中移除。这个过程会迭代地应用于每个注意力算子,最终产生n+1个模型切片,其中n表示原始注意力算子的数量。


图6:一个LLM的分区计算图。

4.2.2 资源利用重叠

注意力计算的拆分。虽然Transformer块中的注意力算子和其他算子是顺序执行的,但仔细研究注意力计算过程会发现实现部分资源利用重叠的潜力。给定一个注意力查询q和词元索引集I,注意力计算可以以分治的方式进行。假设I可以写成两个不相交子集I1和I2的并集,并令


其中 $A_q(I)$ 是注意力输出,$S_q(I)$ 是softmax的分母,那么 $A_q(I)$ 可以通过组合在I1和I2上的部分注意力结果,即 $[A_q(I_1), S_q(I_1)]$ 和 $[A_q(I_2), S_q(I_2)]$ 来轻松获得:

通过提前计算和传输实现重叠。在LLM解码过程中,我们可以在注意力计算期间将当前的词元集分为两个部分:所有先前的词元(prev)和新生成的词元(new)。请注意,一旦 $q_n$ 准备好,$[A_q(prev), S_q(prev)]$ 就可以计算了;因此,我们可以提前执行Q-Proj并传输 $q_n$,然后执行K-Proj、V-Proj并传输 $k_n, v_n$ 到注意力工作节点。如图7所示,这不仅提高了两类工作节点的GPU利用率,还通过将通信隐藏在计算之后来减少了端到端延迟。


图7:通过拆分注意力计算实现资源利用重叠的示意图。

自动化转换器集成。上述注意力拆分优化已集成到我们的自动化模型转换器中。在剖析原始模型后,转换器通过计算其计算图的拓扑排序来为每个模型切片生成一个串行程序。在此拓扑排序过程中,我们总是将Q-Proj算子及其所有依赖项尽可能地提前。然后,我们紧接着在Q-Proj算子之后插入“发送Q”指令,并在该切片的末尾插入“发送KV”指令。

4.3 执行流水线

资源闲置问题。由于基于Transformer的模型的串行特性,如果只有一个批次在处理中,那么当计算设备工作时,内存设备是空闲的,反之亦然。为了解决这种资源利用不足的问题并提高系统吞吐量,我们可以以流水线方式并发运行多个批次。通过精心设计的流水线,可以在不牺牲延迟的情况下实现更高的硬件利用率。我们提出了旋转式交错流水线(rotational staggered pipelining)来解决这个问题。

旋转式交错流水线。假设我们并发执行n个批次。令 $t_m, t_a$ 分别代表执行一个模型切片和一个注意力算子所需的时间。如图8所示,我们部署n-1个模型副本,每个副本的任务启动时间比前一个晚 $t_m \over {n-1}$。所有批次共享一组公共的内存设备,以最大化聚合内存带宽并提高内存利用率。对于每个批次,KV缓存被均匀地划分到这些设备上。所有内存设备共同为一个批次计算注意力算子。内存设备的数量被选择以使 $t_a = {t_m \over {n-1}}$。在注意力算子之后,每个批次根据一个旋转调度转换到下一个模型副本;也就是说,第j个批次的第k个模型切片在副本 $(j + k) \mod (n - 1) + 1$上执行。

流水线优势与挑战。这种旋转式任务调度与交错的执行间隔相结合,保证了每个批次的无缝任务转换,并确保了每个设备上的工作流无冲突、无气泡。此外,通过增加并发批次的数量,由于注意力计算时间的减少,可以降低整体推理延迟。然而,旋转调度需要在计算设备之间迁移批次执行上下文。注意,当n=2时,上下文迁移是不必要的,因为两个批次都在单个模型副本内执行。


图8:旋转式交错流水线示意图。

5. 系统实现

代码与依赖。Lamina用约6000行Python和C/C++代码实现,外加几行实现自定义内核的CUDA代码。全主机旁路网络栈建立在一个修改版的rdma-core【索引6,RDMA core userspace libraries and daemons,https://github.com/linux-rdma/rdma-core】之上。Lamina使用Ray【索引5,Ray,https://www.ray.io/】来促进在分布式异构环境中的任务调度和工作节点放置。

容错机制。在注意力卸载架构中,我们有两种不同类型的加速器。Lamina对这两种加速器的故障采用不同的处理方法。注意,所有请求状态,即KV缓存,仅存储在注意力设备中。因此,如果任何模型工作节点发生故障,我们可以无缝地用一个正常工作的节点替换该故障节点,而不会丢失任何进度。如果注意力工作节点发生故障,我们通过使用存储在LLM服务前端的提示文本和已生成的词元来重建KV缓存。

处理预填充-解码转换。在预填充(prefill)阶段,生成的KV缓存应传输到注意力工作节点以进行解码。对于每个请求,全局调度器选择一组模型工作节点和注意力工作节点来处理解码阶段。与之前的工作【索引40,Splitwise: Efficient generative llm inference using phase splitting,2024;索引59,DistServe: Disaggregating prefill and decoding for goodputoptimized large language model serving,2024,OSDI】类似,KV缓存以逐层异步的方式传输,以将通信延迟隐藏在计算之后。此外,数据传输由注意力工作节点控制:注意力工作节点仅在从模型工作节点接收QKV张量之间的空闲时段从预填充工作节点读取KV缓存。这种方法最大限度地减少了对正在进行的解码任务的干扰。

注意力并行化。鉴于单个设备的能力有限,我们可以使用多个内存设备来共同存储KV缓存并计算注意力算子。如图9所示,注意力算子可以在内存设备之间以多种方式并行化。一种方法是将不同的请求分配到不同的设备上;另一种策略是划分并分配可以独立计算的注意力头到不同的设备。头级分区(head-level partitioning)方法确保了均衡的工作负载分配,而请求级分区(request-level partitioning)可能由于请求之间序列长度和因此KV缓存大小的差异而导致负载不均衡。然而,头级分区的灵活性有限,因为它要求内存设备的数量能被注意力头的数量整除。我们在Lamina中选择了头级分区,因为它能提供最佳的负载均衡。


图9:注意力算子的工作分区方法。

A4 实验环境与结果

实验环境

  • 硬件配置:在一个真实的异构集群上部署Lamina,该集群包含两种类型的GPU节点。每种节点分别包含八个H100或H20 GPU。每个GPU通过PCIe交换机与一个专用的ConnectX-7 NIC配对。GPU节点之间通过400 Gbps RoCE网络互连。实验中,Lamina使用H100作为计算优化型GPU,H20作为内存优化型GPU。
  • 模型:Lamina支持多种LLM架构,包括OPT、LLaMA和LLaMA3。实验选取了LLaMA-33B、LLaMA-65B和LLaMA3-70B进行评估(见表3)。所有模型参数和KV缓存均以FP16格式存储。
    表3:用于评估的大语言模型。
  • 数据集/工作负载:使用从两个LLM服务提供商Azure和Kimi的生产系统中收集的四个请求轨迹(traces)来模拟真实世界的LLM使用场景。这些轨迹只包含用户请求的序列长度。因此,实验中使用具有相同序列长度的虚拟词元请求进行评估。轨迹的摘要信息见表4。
    表4:用于评估的请求轨迹。
  • 软件配置与基线系统:与vLLM【索引28,Efficient memory management for large language model serving with pagedattention,2023,SOSP】进行比较,这是一个为高吞吐量优化的SOTA LLM服务系统。vLLM在同构H100 GPU上运行,并使用张量并行进行多GPU推理。为了公平比较,对vLLM进行了修改以在评估中移除预填充阶段,因为Lamina仅关注解码阶段。

实验结果

6.1 服务性能

等成本配置下的吞吐量对比:在硬件成本相似的配置下(见表5),Lamina与vLLM进行了比较。Lamina用更便宜但内存容量和带宽更大的H20替换了一半的H100。如图10所示,在所有模型和轨迹上,Lamina的吞吐量始终比vLLM高出16.1%至90.1%。这主要归功于Lamina能够支持更大的批次大小,平均是vLLM的2.39倍。这证明Lamina有效利用了内存优化设备提供的额外内存容量来提升解码吞吞吐。
表5:用于评估的等成本硬件配置。



图10:使用大致相等成本硬件的Lamina和vLLM的LLM解码性能指标。

延迟分析:Lamina的词元生成延迟比vLLM高,原因有二:一是Lamina采用了更大的批次,导致模型和注意力工作节点的执行时间更长;二是模型和注意力的解耦可能引入额外的调度和网络开销。尽管如此,Lamina的端到端延迟在大多数情况下仍能满足交互式在线LLM服务的SLO要求。

不同硬件配置下的性能:图11展示了在不同硬件配置下Lamina和vLLM的解码吞吐量。Lamina的吞吐量随着注意力工作节点的增加而迅速提升,因为这允许了更大的批次。增加昂贵的模型工作节点只能通过减少模型部分的执行延迟来温和地提高吞吐量。LLaMA3-70B是个例外,当DOP=(2,4)时,其批次大小已达到800,饱和了模型工作节点的计算资源,因此增加更多内存设备不会显著提升吞吐量。这表明模型和注意力工作节点之间的最佳比例因模型和工作负载而异。


图11:不同硬件配置下的解码吞吐量和硬件成本。图中注释了Lamina的DOP和vLLM的张量并行度。成本效益最佳的配置已加粗显示。

6.2 延迟分解

延迟构成分析:本节测量了不同系统配置下的词元间隔时间(TBT)以及模型、注意力工作节点的执行时间和网络开销。如图12所示,对于较小的批次,模型执行时间主导了词元生成延迟。随着批次增大,注意力和网络延迟迅速增加,而模型执行时间几乎保持不变,这表明计算资源利用率随着批次大小的增加而提高。由于自动化资源利用重叠优化,观察到的TBT可能小于各部分时间之和。


图12:词元生成延迟分解。

6.3 网络栈优化

FHBN性能评估:通过微基准测试评估了全主机旁路网络(FHBN)栈的有效性。在两个不同节点上的GPU之间进行了ping-pong测试。如图13所示,对于小数据量,延迟主要由网络延迟决定,FHBN实现了33.0 µs的端到端延迟,相比NCCL的66.6 µs减少了50.5%。对于大数据量,FHBN的峰值网络带宽达到45.7 GB/s,是线路速率的91.4%,而NCCL仅为35.5 GB/s。这证明了FHBN在数据中心网络中作为现有通信库的优越替代方案。


图13:在两个不同节点上的GPU之间进行网络ping-pong测试,通过400Gbps RoCE互连。

6.4 资源利用重叠

重叠优化效果评估:为了评估自动化模型转换器中实现的资源利用重叠的有效性,对LLaMA-65B和LLaMA3-70B模型进行了开启和关闭该优化的实验。如图14所示,对于LLaMA-65B模型,该优化带来了显著的性能提升,最高可达13.2%,尤其是在批次较大时效果更明显。而对于采用GQA的LLaMA3-70B模型,由于其KV张量尺寸小了8倍,优化的空间较小,最大延迟减少仅为3.5%。


图14:启用和禁用自动资源利用重叠时的词元间隔时间(TBT)结果。

A5 结论

本文提出了一种名为模型-注意力解耦的创新架构方法,以提高LLM解码的效率。该方法源于一个观察:LLM解码阶段可以分为计算密集型部分和内存密集型部分(即注意力算子)。因此,我们可以分别为每个部分使用计算优化和内存优化的设备,以提高硬件资源利用率。为了实现这一想法,我们设计了一个经过改造的、延迟优化的网络栈,以促进远程GPU之间频繁的数据传输。我们还开发了自动化工具,用于转换和优化现有LLM以适应模型-注意力解耦。我们在一个包含异构GPU的集群上开发并部署了Lamina。在从生产系统收集的轨迹上进行的评估表明,在硬件成本相近的情况下,Lamina的吞吐量比异构解决方案高出16.1%至90.1%。

方法细节中引用的参考文献汇总

  • [3] GPUDirect RDMA (官方文档, https://docs.nvidia.com/cuda/gpudirect-rdma/)

    • 引用位置: 第4.1节,描述传统GPU网络通信开销时。
    • 引用原文描述: “To reduce GPU-aware networking overhead, GPUDirect RDMA (GDR) [3] is developed to allow the RDMA-capable NIC (RNIC) to directly access GPU memory.” (为减少GPU感知网络的开销,开发了GPUDirect RDMA (GDR) [3],允许支持RDMA的网卡直接访问GPU内存。)
  • [4] Mellanox adapters programmer’s reference manual (技术手册, https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf)

    • 引用位置: 第4.1节,描述FHBN send实现时。
    • 引用原文描述: “To further decrease latency, we leverage the BlueFlame mechanism, a hardware feature provided by Mellanox RNICs [4].” (为进一步降低延迟,我们利用了BlueFlame机制,这是Mellanox RNIC提供的一个硬件特性 [4].)
  • [5] Ray (开源项目, https://www.ray.io/)

    • 引用位置: 第5节,描述Lamina的实现细节时。
    • 引用原文描述: “Lamina uses Ray [5] to facilitate task scheduling and worker placement in distributed heterogeneous environments.” (Lamina使用Ray [5]来辅助在分布式异构环境中的任务调度和工作节点部署。)
  • [6] RDMA core userspace libraries and daemons (开源库, https://github.com/linux-rdma/rdma-core)

    • 引用位置: 第5节,描述Lamina的实现细节时。
    • 引用原文描述: “The fully host-bypassed network stack is built on top of a modified version of rdma-core [6].” (全主机旁路网络栈是基于一个修改版的rdma-core [6]构建的。)
  • [28] Efficient memory management for large language model serving with pagedattention (Kwon et al., 2023, SOSP)

    • 引用位置: 第5节,作为基线系统的描述。
    • 引用原文描述: “We compare with vLLM [28], a state-ofthe-art LLM serving system optimized for high throughput.” (我们与vLLM [28]进行比较,这是一个为高吞吐量优化的SOTA LLM服务系统。)
  • [40] Splitwise: Efficient generative llm inference using phase splitting (Patel et al., 2024)

    • 引用位置: 第5节,描述处理预填充-解码转换时。
    • 引用原文描述: “Like previous works [40, 59], the KV cache is asynchronously transferred in a layer-by-layer fashion to hide the communication latency behind computation.” (与之前的工作 [40, 59] 类似,KV缓存以逐层异步的方式传输,以将通信延迟隐藏在计算之后。)
  • [57] Orca: A distributed serving system for Transformer-Based generative models (Yu et al., 2022, OSDI)

    • 引用位置: 第5节,描述基线系统vLLM集成的优化时。
    • 引用原文描述: “vLLM also integrates optimizations from other LLM inference systems, such as continuous batching from Orca [57].” (vLLM也集成了来自其他LLM推理系统的优化,例如来自Orca [57]的连续批处理。)
  • [59] DistServe: Disaggregating prefill and decoding for goodputoptimized large language model serving (Zhong et al., 2024, OSDI)

    • 引用位置: 第5节,描述处理预填充-解码转换时。
    • 引用原文描述: “Like previous works [40, 59], the KV cache is asynchronously transferred in a layer-by-layer fashion to hide the communication latency behind computation.” (与之前的工作 [40, 59] 类似,KV缓存以逐层异步的方式传输,以将通信延迟隐藏在计算之后。)