Beluga: A CXL-Based Memory Architecture for Scalable and Efficient LLM KVCache Management
Beluga: A CXL-Based Memory Architecture for Scalable and Efficient LLM KVCache Management
作者/机构: Xinjun Yang, Qingda Hu, Junru Li, Feifei Li, Yicong Zhu, Yuqi Zhou, Qiuru Lin, Jian Dai, Yang Kong, Jiayu Zhang, Guoqiang Xu, Qiang Liu (Alibaba Cloud Computing)
A1 主要贡献
随着大型语言模型(LLM)尺寸的急剧增大和对长上下文推理需求的增长,内存已成为GPU加速的LLM服务中的关键瓶颈。GPU上的高带宽内存(HBM)虽然访问速度快,但容量有限,必须依赖主机内存(CPU DRAM)来支持大规模的键值缓存(KVCache)。然而,DRAM的最大容量受到每个CPU插槽内存通道数量的限制。为了克服这一限制,现有系统通常采用基于RDMA的解耦内存池,但这引入了高访问延迟、复杂的通信协议和同步开销等重大挑战。新兴的CXL技术为KVCache设计带来了新机遇。
本文提出了Beluga,一个新颖的内存架构,它使GPU和CPU能够通过CXL交换机访问一个共享的大规模内存池。该设计支持通过CXL fabric进行原生的加载/存储(load/store)访问语义,从而提供接近本地内存的延迟,同时降低了编程复杂性并最小化了同步开-销。
核心问题:
- 内存瓶颈:LLM推理中庞大的KVCache对内存容量提出了巨大挑战,超出了单机GPU HBM和CPU DRAM的承载能力。
- RDMA的局限性:基于RDMA的内存池方案虽然扩展了容量,但其作为网络协议而非内存总线的设计,带来了额外的 数据移动、高延迟、复杂的编程模型和昂贵的同步开销。
研究目标:
- 设计一个基于新兴CXL技术的内存架构,以克服RDMA方案的局限性。
- 实现一个能够让GPU直接、高效地访问大规模共享内存池的系统。
- 验证该架构在LLM KVCache管理场景下的有效性,提升推理性能。
主要贡献:
- 提出Beluga系统:我们提出了Beluga,一个利用CXL 2.0交换机为GPU集群提供对解耦共享内存池访问的内存系统。据我们所知,Beluga是第一个将GPU集群与CXL交换基础设施集成的系统,展示了通过CXL互连扩展GPU集群内存容量的可行性。
- 全面的性能评估与优化:我们对Beluga在GPU集群上通过CXL 2.0交换机访问内存的性能进行了全面的评估,并提出了一套系统级优化和指导方针以解决关键性能瓶颈。这些优化使Beluga能够有效克服多主机CXL 2.0环境中固有的缓存一致性限制。与基于RDMA的解耦内存系统相比,Beluga在写操作上实现了高达7.0倍的延迟降低,在读操作上实现了高达6.3倍的延迟降低。此外,Beluga提供了可扩展的内存带宽,适用于高性能、数据密集型的GPU工作负载。
- 提出Beluga-KVCache:我们基于CXL内存池构建了一个名为Beluga-KVCache的KVCache管理系统。我们通过将其集成到开源推理库vLLM中,展示了其在高效KVCache读写、基于CXL的RPC和简化调度器方面的有效性。评估表明,与最先进的基于RDMA的解决方案MoonCake相比,Beluga-KVCache将LLM推理吞吐量提高了高达4.79倍。
A3 背景知识
2.1 LLM推理和KVCache
LLM推理过程:大型语言模型(LLM)基于Transformer架构,其推理过程通常包括两个阶段:prefill(预填充)和decode(解码)。计算密集型的prefill阶段并行处理输入提示以生成第一个输出token。随后,decode阶段自回归地生成输出序列,每个新token的生成都依赖于所有先前的token及其对应的KV激活值。
KVCache的作用与需求:KVCache存储了先前token的中间KV激活值,是一种关键的“空间换时间”优化,避免了冗余计算,对于实现低延迟的token生成至关重要。此技术广泛用于多轮对话、共享固定上下文(如系统提示)和RAG中处理长文档等场景。这些场景催生了对高效存储系统的设计需求,以便卸载、重用和共享KVCache。这类系统有四个基本要求:
- 可扩展的容量:KVCache的内存占用随上下文长度线性增长,单个请求就可能消耗GB级的内存。因此,解耦系统必须提供一个能够超越单节点内存限制的大型弹性内存池。
- 高效共享:为了优化多服务器GPU集群的资源利用率和吞吐量,需要一个解耦架构,使所有服务器都能访问统一的KVCache内存池。
- 低延迟访问:使用缓存前缀的主要好处是消除重计算开销,直接改善首token生成时间(TTFT)。因此,从内存池检索KVCache的端到端延迟必须远低于其重计算延迟。
- 高聚合吞吐量:多GPU服务器的内在并行性对内存系统提出了高吞吐量要求。聚合带宽不足会导致GPU停顿,直接限制系统吞吐量,甚至可能使重计算成为更高效的选择。
2.2 远程直接内存访问(RDMA)
RDMA基本原理:RDMA允许网络接口控制器(NIC)直接访问远程主机的内存,绕过远程CPU和操作系统内核。通过最小化网络协议栈开销,RDMA实现了高带宽和低CPU利用率,成为现代数据中心广泛采用的互连技术。
核心原语:大多数RDMA内存池依赖单边通信原语,这有助于节点间直接进行内存到内存的数据传输,消除了内核上下文切换和数据复制等传统开销。关键原语包括用于直接数据检索和放置的RDMA Read和RDMA Write,以及用于在远程内存上执行原子读-改-写操作的RDMA原子操作(如CAS, FAA)。RDMA的单边原语提供了超低延迟、高吞吐量的互连,是构建高性能内存池的基础。
2.3 Compute Express Link(CXL)
CXL协议概述:CXL是一种新兴的互连协议,旨在促进主机CPU和加速器之间的高效内存共享。它定义了三个不同的协议:
- http://CXL.IO:专用于CXL设备管理和DMA数据传输。
- CXL.cache:用于设备对主机内存的一致性访问。
- CXL.mem:使主机能够通过加载/存储(load/store)方式访问设备附加的内存,以实现内存扩展和池化。
CXL标准演进:CXL标准迅速发展,从支持单设备连接的CXL 1.1,到基于交换机的内存池化(CXL 2.0),再到最近支持多层交换和缓存一致性的CXL 3.0,从而实现了更灵活、更强大的异构系统。
本文焦点:本文主要关注利用CXL.mem接口的内存扩展解决方案。目前,支持CXL.mem的硬件设计和生产仍处于早期阶段。现有硬件大多只支持CXL 1.1。尽管一些研究工作开发了支持CXL 2.0的FPGA原型,但其设备数量和内存容量有限。XConn最近推出了首款商用CXL 2.0交换机,该交换机提供CXL.mem接口,支持256个通道和多主机并发访问,最小64字节I/O延迟约为750ns,最大交换容量为2TB/s。CXL 2.0交换机的出现从根本上将该协议提升到超越单节点内存扩展的水平,从而能够创建大规模的共享内存池。
A3 关键Observation/设计原则
3.1 RDMA内存池回顾
传统架构:传统的KVCache管理方法如图2a所示,采用基于RDMA的内存池来聚合集群中的DRAM,以存储卸载的KVCache。这种架构提供了可扩展的内存容量和跨服务器的数据共享。
硬件拓扑:一个典型的多GPU集群通过RDMA网络互连。每个服务器包含主机CPU、多个GPU和RDMA NIC,这些组件通过PCIe交换机互连。例如,一个8-GPU服务器使用四个PCIe交换机,每个交换机连接两个GPU和一个RDMA NIC,并有上行链路到CPU插槽。专用的RDMA NIC提供高达1.6 Tbps的服务器间带宽,使得每台服务器可以贡献其本地DRAM(例如2TB)形成一个分布式共享内存池。
两种访问方式:
- CPU驱动的RDMA内存池访问:vLLM、MoonCake和LMCache等框架采用此方法。主机CPU直接向NIC发出RDMA命令来获取远程数据。存储KVCache的过程分两步:首先,数据从GPU复制到主机内存的“反弹缓冲区”(bounce buffer);然后,CPU启动RDMA传输将数据从该缓冲区移到远程池。加载数据则路径相反。
- GPU驱动的RDMA内存池访问:此方法使用GPUDirect RDMA(GDR),允许GPU直接向NIC发出命令,完全绕过主机CPU。一个专用的GPU内核管理所有网络通信,通过发出RDMA Write/Read并轮询工作完成(WC)信号来存取KVCache。GDR是NVIDIA A100/H100等昂贵的数据中心级GPU的专属功能,在RTX 4090等消费级GPU上不可用。
3.2 RDMA内存池的低效性
性能低效:尽管RDMA提供高吞吐量互连,但它在内存解耦方面存在根本的架构权衡,导致性能低下和系统复杂性。
- 间接的主机中转数据路径:CPU驱动的模型通过主机内存中的“反弹缓冲区”强制所有数据中转,创建了一条昂贵的数据路径(写操作:GPU → 主机DRAM → 远程内存;读操作则相反),这为GPU和内存池之间的数据传输引入了显著的延迟。
- 复杂低效的控制路径:RDMA卸载的一个基本限制是不同组件间需要昂贵的同步。CPU驱动模型需要CPU-GPU协调,GPU驱动模型需要轮询SM和计算SM之间的同步。这种跨组件协调带来了显著的延迟惩罚。我们的微基准测试显示,在H20 GPU上,一个16KB的传输总共耗时10.55µs,但实际数据移动仅占2.68µs,其余近75%(约8µs)是同步开销(如内核启动或等待完成)。此外,KVCache的读写由于GPU和内存池中数据布局不同而变得复杂,例如Qwen-32B(GQA)中的一个KVCache块需要128次非连续的20KB数据传输,这增加了控制路径的额外开销。
- 资源利用效率低下:RDMA基于轮询的控制路径本质上导致资源浪费。在CPU驱动模型中,一个主机线程会消耗整个CPU核心来轮询网络完成事件。在GPU驱动模型中,问题更严重,一个专用的轮询内核会占用宝贵的SM资源,与推理任务产生直接竞争,并使流同步复杂化。
系统复杂性:
- 开发复杂性:RDMA本身编程复杂,开发者需要处理底层网络管理而非简单的内存操作。GPU驱动模型更进一步,要求在CUDA内管理RDMA协议栈并协调轮询和计算流。
- 分层内存的优化复杂性:RDMA访问的高昂开销可能抵消KVCache重用带来的计算节省。开发者必须手动实现请求批处理等关键优化,并通过管理QP排序语义和软件检查来确保单边RDMA的数据一致性。这迫使系统采用复杂的、缓存感知的调度策略,将调度决策与KVCache的局部性紧密耦合,导致调度复杂性高、负载不均和维护开销增加。
A2 方法细节
4. Beluga:一个基于CXL的内存架构
为了克服RDMA的局限性,我们提出了Beluga架构,它利用业界首款可量产的CXL交换机(XConn XC50256)来构建一个可扩展的共享内存池。这使得GPU能够通过简单的加载/存储操作访问池化内存,解决了RDMA方法的性能瓶颈和编程复杂性。
4.1 概述
架构设计:如图2b所示,Beluga架构用两个PCIe/CXL适配器替换了四个专用的RDMA NIC。每个服务器有两个CPU插槽(NUMA架构),每个插槽通过一个PCIe 5.0 x16的PCIe/CXL适配器连接到CXL交换机。CXL内存池本身由一个交换机节点和一个独立的内存盒组成。
核心硬件:交换机核心配备了两颗XConn XC50256芯片,每颗芯片通过256个PCIe 5.0通道提供2TB/s的转发能力。这些通道通常在CXL内存设备和计算服务器之间平均分配。底层的CXL交换机最多可将16台服务器连接到一个8TB的内存池,总带宽为1TB/s。这种连接性使Beluga能通过其内部地址映射和转发逻辑支持并发的多主机访问。
4.2 Beluga的内在优势
内存访问范式转变:通过将内存访问范式从网络协议(RDMA)转变为内存语义接口(CXL),Beluga相比基于RDMA的方法提供了几个直接的优势。
性能提升:如图4所示,Beluga为CPU和GPU提供了标准的数据访问接口。
- CPU访问:支持 (1) 通过加载/存储指令直接内存访问;(2) 使用Intel数据流加速器(DSA)进行硬件加速传输。
- GPU访问:支持 (1) 通过cudaMemcpy API进行直接的点对点(P2P)传输;(2) 通过自定义CUDA内核进行细粒度的非连续访问。
- 数据路径优势:Beluga允许GPU直接访问全局内存池,消除了CPU驱动RDMA所需的多步数据路径和反弹缓冲区,显著降低了延迟。
- 控制路径优势:Beluga中的数据传输内核无缝集成到GPU原生的CUDA流中。这种统一的控制消除了RDMA中固有的昂贵的跨组件同步,从而移除了外部CPU协调或内部GPU轮询相关的开销。
系统简化:
- 编程模型:数据访问接口类似于本地DRAM的接口。开发者无需管理底层网络协议栈、复杂的工作请求准备以及RDMA内存池所需的性能优化。消除GPU的跨组件同步开销也简化了数据传输和计算之间的协调。
- 内存管理:Beluga提供了一个统一的地址空间和简化的空间控制。启动时,每个主机的BIOS会检测并为连接的CXL设备保留一块连续的物理地址空间。Beluga利用此基础,让每个主机以直接访问(DAX)模式管理CXL内存池。在此模式下,CXL内存被暴露为块设备,允许用户空间进程使用mmap()将整个内存区域直接映射到其虚拟地址空间。这种直接映射机制是资源划分和数据共享的关键。应用程序可以为不同主机分配不同的内存偏移来实现逻辑资源划分,或让多个主机映射相同区域以实现高效数据共享。
- 硬件成本:高速网络硬件(如400Gbps NIC)对于LLM推理的典型带宽需求通常是过度配置的。用更具成本效益的CXL组件替换昂贵的RDMA NIC可以显著降低总拥有成本。如表1所示,CXL方法降低了主机接口卡和交换机的设备级成本。对于内存设备,与主机中的传统DIMM内存不同,CXL内存设备可以使用密度较低的芯片,降低每GB的成本。在系统层面,CXL内存池将内存与CPU资源分离,允许多台服务器共享内存,提高云环境中的利用率。
总结:对于单机架内存池,CXL通过在性能、成本效益和资源利用率之间提供最佳平衡,显示出比RDMA明显的优势。
Table 1: 硬件成本分析。
B1 样品价格,最终价格可能会有变动。
5. Beluga的特性分析与优化
尽管CXL提供了类似内存的简单接口,但其CXL 2.0硬件的性能,特别是对于GPU直接到内存的访问,尚不明确,这阻碍了软件发挥其全部潜力。本节对Beluga中所有数据路径进行了全面的性能分析,并提出了一系列实用优化,总结于表3。分析分为三部分:首先,设计和评估三种基于软件的一致性方法(§5.1);其次,分析不同I/O传输方法的延迟,并为延迟敏感型工作负载确定最佳方法(§5.2);最后,评估并确定Beluga中的带宽瓶颈,并提出两种有效的解决方案来扩展带宽密集型工作负载的性能(§5.3)。所有实验均在商用硬件平台上验证,详细规格见表2。
Table 2: 实验设置。
5.1 在非一致性CXL上进行数据共享
目标和方法:虽然CXL 2.0中的CXL.cache和CXL.mem支持主机CPU和设备之间的内存共享一致性,但它们不支持跨多个主机CPU的缓存一致性。CXL交换机向主机处理器呈现一个逻辑上统一的内存地址空间,但每个计算节点维护一个独立的缓存层次结构(L1/L2/L3),与其他节点隔离。因此,一个主机的写操作仅限于其本地缓存,不会自动传播到远程主机。在KVCache场景中,这会导致不准确、不一致或不连贯的输出。为了解决这个问题,需要软件层面的显式跨节点同步协议来协调所有缓存状态转换。我们的系统专注于优化KVCache共享机制,支持单写入者插入KVCache块,并允许多个读取者访问,从而消除冗余计算。我们为此场景提出了三种写入端和两种读取端的一致性管理方法,并进行实验以测量延迟并选择最佳方法。
写入端:确保数据到达CXL内存:为确保一个主机的写操作对其他主机可见,数据必须从其私有缓存层次结构刷新到CXL内存。有三种方法可以实现:
- 静态不可缓存配置:通过在内存类型范围寄存器(MTRRs)中设置相应的内存类型,将CXL内存区域配置为不可缓存。这样,CPU发起的写操作(标准存储指令或DSA传输)会绕过主机缓存,直接发送到Beluga CXL内存池。对于从GPU进行的设备到主机(D2H)memcpy,我们禁用数据直接I/O(DDIO),强制D2H memcpy流量绕过末级缓存(LLC),直接写入CXL内存池。
- 写后细粒度缓存刷新:此方法允许缓存数据,但依赖软件使用CPU指令(如CLFLUSH、CLFLUSHOPT和CLWB)显式地将修改后的缓存行刷新到内存。然而,这些方法会带来很高的指令开销,因为它们是按缓存行操作的,开销与数据大小成线性关系。
- 细粒度绕过缓存写入(无缓存刷新):为消除显式刷新的开销,可以利用CPU中一些绕过缓存的指令。CPU发起的写操作可以使用非临时存储指令(ntstore),它绕过缓存层次结构直接将数据写入CXL内存池。同样,Intel DSA提供一个缓存绕过标志,使DMA操作能够将其有效负载直接写入内存,而不填充CPU缓存。
读取端:确保从CXL内存读取新数据:同样,为了防止读取本地缓存中的陈旧数据,系统应将内存区域配置为不可缓存,或在读取操作前显式地使目标缓存行无效。
- 静态不可缓存配置:对于CPU读取,不可缓存属性通过强制每次加载都访问远程CXL内存池来确保正确性。即使对于设备发起的传输(如cudaMemcpy),GPU对CXL内存的正确访问也需要不可缓存属性,因为并发或之前的CPU活动可能已创建数据的缓存副本。
- 读前细粒度缓存刷新:此方法通过在读取操作前使用CLFLUSH指令主动使陈旧缓存行无效来确保读取的正确性。CLWB不适用于此目的。这些刷新指令由于按缓存行操作而开销较高。
与RDMA的对比:上述写入或读取端的设计很简单:设置缓存状态和禁用DDIO是一次性步骤,缓存刷新和读/写一样简单。相比之下,在RDMA中实现缓存一致性需要更复杂的异步完成和队列对(QP)排序管理。
(Exp #1) 性能特性分析:我们评估了16 KB读/写的延迟,结果如表4所示。对于写操作,非临时存储(ntstore)是CPU发起存储的最有效方法,延迟为2.41µs。使用不可缓存内存则非常慢(281.56µs)。对于DSA,使用不可缓存内存和缓存绕过标志的性能几乎相同且最佳。对于GPU D2H传输,禁用DDIO以绕过缓存(9.14µs)比依赖后续CPU刷新(11.06µs)更有效。对于读操作,CPU从不可缓存内存加载非常慢(166.49µs),因此唯一可行的方法是在读取前刷新缓存(5.98µs)。DSA读取和GPU H2D复制在直接从不可缓存内存读取时性能更好。
优化:由于当前CXL 2.0交换机不支持主机间的硬件缓存一致性,需要软件来管理。最佳性能的实现依赖于根据操作发起者采用特定的缓存管理策略。我们总结出三个设计提示:
- O1:对于CPU的存储/加载,写操作应使用非临时存储(ntstore),读操作前应执行CLFLUSH。
- O2:对于Intel DSA引擎,对不可缓存内存进行操作是读写的最佳实践。
- O3:对于GPU memcpy操作,禁用DDIO并使用不可缓存内存可获得最高效率。
未来,CXL 3.0交换机可能通过硬件提供跨主机的缓存一致性,从而无需软件同步。但CXL 3.0的一致性仍在发展中,我们的软件设计对于非一致性区域仍然是必要的,并且可以与未来的硬件一致性区域结合使用以进一步简化和加速。
Table 3: Beluga中的性能优化。
Table 4: 缓存一致性方法的延迟(µs)。粗体条目突出显示了每个操作的最佳机制。(实验 #1)
*自定义内核的结果包括内核启动时间。
5.2 延迟优化
目标和方法:CXL内存访问的性能严重依赖于具体操作和访问大小。我们首先通过特性分析Beluga中四种基本操作的延迟,以确定它们的最佳使用场景。此外,我们将Beluga与本地内存和基于RDMA的内存池进行比较,以突出其性能优势最显著的场景。
(Exp #2) 性能特性分析:图5展示了CPU/GPU与内存池之间的I/O操作延迟。所有延迟基准测试均使用队列深度一(QD=1)来测量单请求延迟。主要观察如下:
- GPU访问性能:直接从CXL到GPU的数据传输与传统的CPU到GPU路径相比具有很强的竞争力(图5c和5d)。在64KB传输大小时,CXL到GPU的复制延迟为11.73µs,非常接近CPU到GPU的10.32µs。这表明CXL内存可以作为GPU工作负载的主要数据源,而延迟开销可忽略不计。
- CPU与DMA的权衡:在将数据移动到CXL时,基于CPU的复制和基于DMA的传输之间存在明显的性能权衡(图5a和5b)。对于小于4KB的I/O,直接CPU加载/存储操作的延迟低于DMA。对于大于4KB的传输,DMA的并行性优势开始超过其设置开销。性能交叉点在16KB,此时Intel DSA引擎优于基于CPU的复制。
- CUDA Memcpy开销:分析显示,CUDA Memcpy的延迟主要由软件开销(主要是GPU驱动栈中的内核启动延迟)主导,而非物理数据移动时间(图5c和5d)。例如,一个16KB的H2D传输耗时10.55µs,但NVIDIA Nsight Systems分析显示实际数据传输仅需2.68µs。因此,建议使用自定义数据复制内核将多个复制任务合并到一个内核中,以避免内核启动和CPU同步开销,从而获得比CUDA Memcpy和GDR更低的延迟。
- 对不可缓存内存的特殊情况:标准cudaMemcpy在源内存为不可缓存(Uncachable)时,对于小于24KB的H2D传输性能会急剧下降,耗时约1.23ms(图5d中未显示)。我们推测CUDA运行时使用基于CPU的指令来优化小传输,这种策略不适用于不可缓存内存。因此,对于从不可缓存CXL内存到GPU的小于24KB的传输,必须使用自定义CUDA memcpy内核。
优化:Beluga提供的延迟与本地DRAM相当,这使得创建扁平、统一的内存池成为可能,有潜力消除对复杂多层内存层次结构的需求。我们提出以下三个优化建议:
- O4:CPU应使用直接加载/存储处理小I/O(< 4KB),使用DSA处理更大的传输。
- O5:在GPU和CXL内存池之间移动数据时,应将任务合并到单个内核中,以隐藏内核启动延迟。
- O6:对于从CXL内存到GPU的大小不超过24KB的数据传输,建议使用自定义实现的CUDA memcpy内核。
5.3 带宽优化
目标和方法:在使用单个PCIe/CXL适配器(16通道)时,我们在Beluga中观察到两个性能异常。首先,读/写性能不对称:从CPU到CXL内存池的读带宽达到预期的46.2 GB/s,与RDMA相当,但写吞吐量仅为33 GB/s。其次,GPU访问CXL内存池时性能显著下降,吞吐量降至26 GB/s,远低于CXL内存控制器和GPU自身PCIe带宽(55.4 GB/s)。从CPU/GPU到CXL内存池的数据路径复杂,穿过PCIe交换机、根复合体(Root Complex)和CXL交换机等组件,其中任何一个都可能成为性能瓶颈。我们的方法是首先对该路径进行瓶颈分析,然后提出针对性优化。
性能特性分析:通过系统性的实证分析,我们确定CPU的根复合体(Root Complex, RC)是CXL内存池架构中的主要性能瓶颈。为了验证这一结论,我们构建了一个专门的微基准测试(图6),测量连接到不同PCIe交换机的GPU和NIC之间的带宽,强制数据路径通过CPU的根复合体。实验结果与CXL内存池的性能完全匹配:PCIe P2P写(即NIC⇒GPU)带宽为33 GB/s,P2P读(即GPU⇒NIC)带宽限制在23 GB/s。这表明瓶颈可能不在CXL内存设备或CXL链路本身,而是受限于根复合体的底层点对点能力。进一步的多GPU-NIC对实验显示,总带宽随对数线性扩展,表明限制在于根复合体内的每通道或每链路资源,而非整体吞吐量上限。
内存设备瓶颈:此外,我们观察到内存设备本身也可能成为瓶颈。每个设备支持22.5 GB/s的带宽。如果来自服务器的所有工作负载都指向单个设备,带宽将受限于该设备。因此,在Beluga上运行的软件应将数据分布到不同的内存设备上,以避免单设备瓶颈。我们当前的实现依赖于2MB粒度的软件内存交错。即将推出的Intel Granite Rapids处理器将提供对CXL设备交错的硬件支持。
优化:基于我们的观察,我们为带宽敏感场景提出了三个优化策略:
- O7:未来架构应支持GPU到CXL交换机的直接连接。这样Beluga可以绕过RC,避免GPU和CXL内存之间的主要瓶颈。
- O8:单个PCIe/CXL适配器的带宽受限于PCIe 5.0 x16。为获得更高带宽,应根据应用需求增加PCIe/CXL适配器的数量。
- O9:Beluga应在多个CXL内存设备间交错数据。此策略可并行化访问,有助于最大化聚合吞D吐量。
5.4 复杂工作负载下的性能
(Exp #3)倾斜内存访问下的性能:我们在并发内存访问下测量了带宽、中位延迟和p99尾延迟,结果如图7所示。对于本地DRAM和CXL访问,写操作使用ntstore,读操作前使用clflush。实验还对比了启用和禁用2MB内存交错的性能,证明了在倾斜工作负载下启用内存交错的必要性。我们有两个观察:
- 首先,CXL的性能优于RDMA。在中位延迟方面,CXL在64B操作中仅为RDMA延迟的10.2%~13.3%,在16KB操作中为39.5%~56.2%。值得注意的是,CXL在16KB写操作中实现了与本地DRAM相当的写延迟。
- 其次,在16KB倾斜工作负载下,没有内存交错的CXL表现出较低的带宽和较高的延迟。这是因为CXL内存盒中的第一个内存设备成为瓶颈,在并发访问期间引入了排队延迟。
(Exp #4) 背景工作负载的影响:该实验在一个内存设备上施加来自背景服务器(服务器B)的固定带宽压力,同时另一台服务器(服务器A)对同一设备执行64B的读/写操作。图8显示了在不同背景带宽压力下的读/写延迟。结果表明,虽然中位延迟无论背景带宽如何都保持稳定,但当背景带宽压力与操作方向相同时,p99延迟会增加。这也展示了CXL内存池的双向能力。
6. 使用Beluga管理LLM KVCache
本节阐述了如何将Beluga与流行的LLM推理框架(如vLLM)集成,并高效管理KVCache。如图9所示,一个LLM推理系统通常由三个关键组件构成:一个大型共享内存池用于存储KVCache;一个全局索引用于将token块映射到其在池中的物理地址;以及一个集中式调度器用于向不同的LLM实例分派请求。Beluga-KVCache与这三个组件都进行了集成。
6.1 KVCache的数据传输操作
KVCache的内存布局与传输模式:在基于Transformer的LLM推理系统中,每个token的KVCache通常以高度碎片化的内存布局存储。如图10所示,一个token的键(key)和值(value)张量在不同注意力层之间通常非连续存储。此外,在单层内,来自不同token的KV张量也非连续存储。然而,LLM推理系统通常必须将KVCache数据序列化为内存池中的连续块以便有效使用,这导致了显著的KVCache传输开销。KVCache传输有两种访问模式:
- Gather write (KVCache写):将来自多个非连续GPU位置的数据收集到一个远程池中的连续块中。
- Scatter read (KVCache读):将来自远程池中一个连续块的数据分散到多个非连续的GPU位置。
RDMA与Beluga的对比:RDMA依赖散点/聚集列表(sglists)来处理这些传输,但对于KVCache效率低下。例如,一个Qwen-32B模型的单个token的KVCache可能分裂成128个非连续块,而ConnectX-7 NIC的硬件限制(sglists最多30个条目)导致需要复杂的逻辑将操作拆分为多个RDMA请求。当利用KVCache稀疏性时,问题变得更具挑战性,访问模式变为非连续到非连续的读取,可能触发IOPS瓶颈,RDMA同样效率低下。
Beluga的优势:相比之下,Beluga利用细粒度的自定义复制内核来处理无限的gather writes或scatter reads,消除了RDMA中繁琐的请求管理。通过将KVCache传输任务从RDMA转移到CXL内存池,Beluga从根本上解决了LLM推理系统中的KVCache迁移瓶颈。
6.2 基于CXL的RPC
背景与需求:对于具有共享KVCache存储的LLM推理系统,高效的元数据管理至关重要。大多数系统依赖集中式元数据服务来跟踪KVCache块的位置,这会产生频繁的远程过程调用(RPC)。传统上,这些RPC通过RDMA或TCP/IP协议实现,由于网络协议栈开销或内核-用户空间转换而导致高延迟。
Beluga的实现:Beluga通过共享内存实现RPC,充分利用CXL提供的加载/存储语义。我们在Beluga中保留一小块内存空间以实现跨服务器通信。Beluga以生产者-消费者的模式建立元数据服务器和客户端之间的通信。客户端和服务器在CXL内存池中预分配固定大小的内存槽用于请求和响应缓冲。客户端(生产者)将请求写入一个空闲槽并更新状态标志(REQ_READY)。服务器(消费者)通过内部自旋等待循环持续轮询状态标志。在处理请求后,服务器将结果写入专用的回复槽并设置RESP_READY标志。
优化措施:我们应用了四项优化以实现高性能:
1. 客户端采用非临时存储(ntstore)指令,避免CPU缓存污染。
2. 元数据服务器在读取前执行CLFLUSH命令,以确保看到客户端写入的最新数据。
3. 我们对并发RPC请求的mfence操作进行批处理,并确保RPC中数据的缓存行对齐。
4. 整个生产-消费过程通过将所有操作保持在用户空间来消除内核模式转换和上下文切换。
通过这些技术,Beluga能够实现比RDMA或TCP/IP更高效的RPC。
6.3 无需KVCache层次的调度
传统调度策略的弊端:使用基于RDMA内存池的传统LLM推理系统(如NVIDIA的Dynamo和MoonCake)存在本地和远程内存访问之间的显著性能差距。为弥补这一差距,这些系统必须采用精心设计的缓存感知调度策略,通过将请求路由到已托管所需数据块的节点来强制实现KVCache局部性,从而避免昂贵的远程抓取。这种策略导致了调度复杂性、负载不均衡和维护开销等明显缺点。
Beluga的简化调度:Beluga通过利用具有近乎本地访问延迟的CXL内存池,在很大程度上消除了上述限制。整个CXL内存池被抽象为一个统一且对称的地址空间,远程CXL内存访问延迟与本地缓冲区访问延迟相当(如§5.2所示),消除了由RDMA引起的性能差异。我们对比了将KVCache卸载到本地内存与Beluga的性能,结果显示对于缓存命中请求,Beluga的TTFT与本地内存极具竞争力。因此,Beluga的设计支持缓存无关(cache-oblivious)调度。传入的请求可以使用标准的负载均衡技术轻松分发,忽略缓存局部性。此外,LLM推理节点可以随时增减,而无需重新平衡KVCache分区,因为访问开销保持一致。
范式转变:通过将计算调度与KVCache局部性解耦,Beluga显著降低了调度开销,并能在倾斜工作负载下保持高吞吐量。这种范式转变为LLM推理系统优先考虑计算效率而非内存层次结构优化提供了可能。
A4 实验环境
-
数据集:
- 主要负载: LV-Eval [60],一个长上下文问答(QA)数据集,所有输入序列超过15K tokens。
- 变体: 通过限制LV-Eval的输入上下文长度,创建了LV-Eval-2K, LV-Eval-4K, 和 LV-Eval-8K三个变体。
-
模型架构:
- 默认模型: 未量化的Qwen-32B模型。
-
硬件配置:
- 计算平台: 一个由两台服务器组成的集群,每台服务器配备8个H20 GPU(每个96 GB),共计16个并发vLLM实例。
- CPU: 2 x Intel(R) Xeon(R) Platinum 8575C。
- 互联: Beluga系统使用2个PCIe 5.0 x16适配器连接到一个8TB的CXL内存池。RDMA基线系统使用4个ConnectX-7(双端口200Gbps)NIC连接到一个4TB的RDMA内存池。详细配置见表2。
-
软件配置:
- 推理框架: vLLM V1 (v0.8.5),启用了其前缀缓存功能。
-
基线系统:
- MoonCake (v3.2),一个基于vLLM和LMCache (v0.3.1) [17, 26]的高度优化的RDMA方案。
- Dynamo (v0.4.1),一个来自Nvidia的RDMA方案。
-
约束: 为了公平比较,所有方案的内存池缓存容量均限制为2TB。
A4 实验结果
7.1 端到端性能评估
为了区分初始缓存填充和后续缓存重用的性能影响,实验设计了两种场景:
- 缓存填充(首次运行):测量初始缓存填充期间的性能,此时vLLM计算KVCache并存入Beluga。
- 缓存命中(第二次运行):所有KVCache已预先填充在共享内存池中,测量预填充阶段完全由缓存重用加速时的性能。
(Exp #5) 结果:如表5所示,Beluga-KVCache在所有指标上都持续优于基于RDMA的基线。
- 在缓存填充场景(30%缓存命中率)中,Beluga-KVCache相较于MoonCake,平均TTFT降低了12.4%,QPS提高了21.5%。
- 在缓存命中场景中,性能优势最为显著:Beluga-KVCache的平均TTFT降低了89.6%,QPS提升了7.35倍。
MoonCake性能较低的原因在于CXL与RDMA技术间的固有性能差距,以及其实现中关键路径上的额外开销(如额外的内存拷贝和分配)。
Table 5: LV-Eval上的推理性能 [60] (实验 #5)。
7.2 敏感性分析
(Exp #6) 对请求到达率的敏感性:如图11所示,在第二次运行(完全缓存命中)中,随着请求率从0.3 QPS增加到9.0 QPS,Beluga-KVCache在TTFT和TPOT方面始终优于MoonCake。这是因为在此场景下,KVCache读取时间成为主要瓶颈,而Beluga基于CXL的数据访问效率远高于MoonCake的RDMA。
(Exp #7) 对输入上下文长度的敏感性:如图12所示,在使用2K、4K、8K和原始LV-EVAL(>15K)不同长度上下文的负载进行评估时,Beluga的性能优势随着输入token数量的增加而变得更加显著。这是因为在更长的上下文中,KVCache的读/写时间占端到端延迟的比重更大。
(Exp #8) 对软件配置的敏感性:
- Prefill-Decode分离架构:在Prefill和Decode节点分离的架构中,Beluga优化的KVCache加载/存储路径使其QPS比MoonCake高3.41倍至9.47倍(图13a, 13b)。
- KVCache块大小:RDMA方案需要通过批处理来摊销控制开销,因此LMCache默认使用256-token的大块。而Beluga可以高效地使用vLLM原生的16-token小块。如图13c所示,MoonCake在使用16-token块时TTFT急剧增加到76.8秒,甚至超过了首次运行的重计算延迟,而Beluga则无此问题。
- 内存交错:启用软件内存交错后,Beluga在缓存命中场景下的QPS从8.49提升到11.32,提高了33.2%,证明了交错能有效分散访问负载,减少竞争。
7.3 性能提升分解
(Exp #9) 密集KVCache数据传输性能:我们针对KVCache传输中的关键操作——散点/聚集操作进行了基准测试。如图14所示,与MoonCake采用的低效两步数据路径(GPU → 主机 → 内存池)相比,Beluga完全消除了主机内存反弹缓冲区。这种直接数据路径使KVCache的写和读延迟分别降低了36.2%和38.7%。
(Exp #10) 稀疏KVCache数据传输性能:我们分析了Qwen-32B和Llama-3-8B模型中的KVCache稀疏性,发现超过74%的重要token是非连续的,确认了稀疏KVCache的内存访问模式是高度离散的(表6)。基于此,我们评估了读取16个稀疏token的KVCache的性能。在Qwen-32B模型中,Beluga实现了95.9%的延迟降低。原因是RDMA方案因需要为非连续数据发出大量高开销请求而成为瓶颈,而Beluga允许单个CUDA内核通过CXL的直接映射内存高效地管理整个细粒度传输。
(Exp #11) RPC性能:我们比较了CXL-RPC与两种RDMA实现(RDMA-RC和RDMA-UD)的性能。如图15所示,在低并发(QD=1)下,CXL-RPC的往返延迟为2.11µs,比RDMA-RC(8.39µs)和RDMA-UD(8.83µs)快4倍。在高并发(QD=128)下,CXL-RPC的单线程吞吐量达到12.13 Mops,分别比RDMA-RC(4.5 Mops)和RDMA-UD(6.65 Mops)高2.7倍和1.8倍。这表明CXL的直接加载/存储模型避免了RDMA协议中大量的软硬件开销。
Table 6: 稀疏性分析和读取延迟(实验 #10)。
A7 补充细节
8. CXL交换机的未来工作
本文使用首批商用CXL 2.0交换机之一,对CXL交换内存池进行了基础性分析,为CXL 2.0系统建立了关键性能基线,并展示了其在KVCache卸载等复杂任务中的有效性,为未来CXL原生系统开辟了广阔的设计空间。
未来硬件架构:随着CXL 3.1的出现,我们设想未来将构建一个基于CXL fabric的大规模、完全解耦的架构,如图16所示。在该架构中,一个对称的fabric将连接计算、内存和存储资源池,为所有连接设备提供统一的低延迟访问。通过将GPU直接连接到此fabric,可以消除遍历主机PCIe交换机和根复合体的延迟开销,提供比Beluga架构更低的内存访问延迟。
未来软件设计:新的硬件架构需要相应的软件栈来管理资源池和大规模一致性数据共享。
- 资源调度器:需要开发能够跨多个CXL交换机动态管理GPU/CPU资源的调度器。
- 一致性协议:设计更高效的CXL内存池一致性协议是关键研究领域,方向包括:(1) 利用应用层语义放宽一致性要求;(2) 使用交换机内资源实现可扩展的、基于目录的一致性;(3) 开发混合模型,用一个小的硬件一致性区域管理更大地址空间的一致性元数据。
未来数据库设计:CXL内存池为需要大规模随机访问能力的内存密集型数据管理系统(特别是向量数据库和图数据库)带来了新机遇。例如,像HNSW这样的图算法传统上需要完全内存存储,可以利用CXL内存池突破内存容量限制。
9. 相关工作
LLM推理中的数据管理:KVCache是LLM推理的主要内存瓶颈,催生了大量旨在减少其内存占用的研究。算法层面有MQA、GQA等架构创新,以及H2O、FastGen、SnapKV等利用注意力分数稀疏性进行压缩或驱逐的方法。系统层面,PagedAttention引入了虚拟内存机制,而MoonCake和Dynamo利用远程资源池实现大规模请求调度。Beluga在此基础上提出了一种新的硬件架构,从系统层面减少数据移动开销,并为稀疏KVCache等新算法优化提供灵活的随机访问支持。
用于数据管理的解耦内存架构:解耦内存已在数据库系统中广泛应用。基于RDMA的内存解耦已用于事务处理和存储优化等多种场景。最近的系统如MoonCake和Dynamo利用RDMA将KVCache管理卸载到内存池。然而RDMA作为网络协议引入了额外的复杂性和延迟。CXL通过提供原生的类内存接口和特性解决了这些限制。此外,还存在如华为CloudMatrix这样的专用解决方案,但其厂商特定的设计限制了其广泛采用。
基于CXL的内存架构:CXL的引入引发了对其性能特性和系统潜力的广泛研究。早期的特性研究利用FPGA原型或CXL 1.1硬件进行了初步的量化分析。随后,一系列研究展示了CXL在内存数据库、图处理和深度学习系统等多种实际场景中的实用性。多项工作对CXL和RDMA进行了仔细的比较分析。先前的工作主要探讨了CXL用于CPU内存扩展。本文将CXL的能力扩展到以GPU为中心的KVCache管理,并使用了商用的CXL 2.0交换机。
A5 结论
本文设计、评估并优化了Beluga,一个新颖的、基于CXL交换机的GPU集群内存系统。通过对商用CXL硬件的详细特性分析,我们识别了关键的性能特征,并为构建高性能应用提出了一系列优化措施。为了展示Beluga的有效性,我们设计并实现了Beluga-KVCache,一个专为LLM推理中大规模KVCache管理量身定制的系统。该系统包括优化的数据传输、轻量级RPC和基于Beluga的简化调度器。我们的评估表明,Beluga-KVCache的性能显著优于最先进的基于RDMA的解决方案MoonCake,突显了CXL作为未来解耦内存系统基础的巨大潜力。
💬 评论讨论
欢迎在这里分享您的想法和见解!