Splitwise：使用相位分裂实现高效生成式 LLM 推理

24年5月来自华盛顿大学和微软的论文“Splitwise: Efficient Generative LLM Inference Using Phase Splitting”。

生成式大语言模型 (LLM) 应用正在快速增长，导致昂贵且耗电的 GPU 大规模部署。 在 LLM 推理，每个推理请求都会经历两个阶段：计算密集型提示计算阶段和内存密集型 token 生成阶段，每个阶段都具有不同的延迟、吞吐量、内存和功率特性。尽管采用了最先进的批处理和调度，但 token 生成阶段并未充分利用计算资源。与提示计算不同，token 生成不需要最新 GPU 那样的计算能力，并且可以以较低的功率和成本运行。

基于这些见解，Splitwise 作为一种模型部署和调度技术，将 LLM 推理请求的两个阶段拆分到单独的机器上。Splitwise 使用适合每个阶段的硬件来实现特定阶段的资源管理。使用当今 GPU 集群中快速背板互连支持的优化网络库，请求状态可在机器之间高效传输。通过 Splitwise，设计针对吞吐量、成本和功率优化的同构和异构 LLM 推理集群。与当前设计相比，Splitwise 集群的吞吐量提高 1.4 倍，成本降低 20%。或者在相同的功率和成本预算下提供 2.35 倍的吞吐量。

生成式大语言模型 (LLM) 的最新进展显著提高了其响应质量和准确性 [18][71]。这些趋势导致 LLM 在各个领域得到广泛采用 [6][21]。大多数现代 LLM 都用 Transformer 架构构建 [77][78]，并表现出相似的特性 [63]。Transformer 模型的大小稳步增长，从拥有 3.4 亿个参数的早期 BERT 模型 [36]，到拥有惊人的 1750 亿个参数的 GPT-3 [28]，以及传言拥有更多参数的 GPT-4。

LLM 通常在昂贵且耗电的 GPU 上运行 [16]。LLM 的大规模部署导致全球 GPU 容量紧缩 [14]。由于大量应用程序利用 LLM，LLM 推理的计算需求远远超过训练的需求。此外，由于训练 LLM 需要昂贵且专用的超级计算机 [56][60]，因此需要大量推理来摊销高昂的训练成本。LLM 推理作业虽然比训练小几个数量级，但考虑到所涉及的计算量，成本仍然很高。

由于输出 token 是逐个生成的，因此单个请求的生成式 LLM 推理由模型的几次前向传递组成。这本质上有两个对比的计算阶段。首先是提示计算阶段，其中所有输入提示 token 并行地通过模型的前向传递，生成第一个输出 token。此阶段往往是计算密集型的，需要当今最新 GPU 的高 FLOPs（每秒浮点运算次数）。其次是 token 生成阶段，其中后续输出 token 是根据最后一个 token 的前向传递和序列中先前 token 的所有缓存上下文按顺序生成的。由于缺乏计算并行性，尽管采用最先进的批处理技术，此阶段往往受到内存带宽和容量的更多限制。在同一台机器上运行这两个阶段，通常会导致不一致的端到端延迟，因为提示和token阶段的批处理是任意的。由于这些挑战，服务需要过度配置昂贵的 GPU，以满足交互式应用程序严格的推理服务级目标 (SLO)。与此同时，云服务提供商 (CSP) 不得不建立大量新的数据中心来满足 GPU 需求，并且还遇到功率瓶颈 [19]。

业界继续发布新的计算能力强大 GPU，每一款都比上一款耗电更多、价格更昂贵。然而，如表所示，最近这些 GPU 上的高带宽内存 (HBM) 容量和带宽，并没有以相同的速度扩展。与前代 A100 GPU 相比，最新的 NVIDIA H100 GPU 的计算能力提高了 3.43 倍，功率提高了 1.75 倍。然而它们的内存带宽仅增长了1.6倍，内存容量并没有增加。

如图显示了生成式 LLM 推理的示例。收到提示查询后，所有输入tokens都会在一次迭代中并行计算，生成第一个token。称之为提示处理阶段。提示计算期间从注意层生成的上下文保存在 KV 缓存中，因为它是所有未来tokens生成迭代所需的。生成第一个 token 后，后续tokens仅使用最后生成的tokens和 KV 缓存，作为模型前向传递的输入。这使得后续tokens生成，比计算繁重的提示阶段更耗费内存带宽和容量。

先前的研究提出了 LLM 推理的三个主要指标：端到端 (E2E) 延迟、第一个token时间 (TTFT) 和吞吐量。再加另一个延迟指标：token之间的时间 (TBT)，跟踪token序贯生成的在线流式吞吐量。下表总结了所考虑的关键性能指标。

推理请求可以一起批处理以提高吞吐量。先前的一些研究已经探索批处理 [23]、[81]。如图显示了使用三种常见批处理机制进行推理的时间线。默认机制仅在请求级进行批处理（图（a））。在这种情况下，就绪请求被一起批处理，但这些请求的所有前向传递，都在运行任何其他请求之前完成。由于请求可能具有较长的token生成阶段，这可能导致请求在期间到达要等待很长时间，从而导致高 TTFT 和高 E2E 延迟。一种优化是连续批处理 [81]（图 （b））。在这种情况下，在模型的每次前向传递之前做出调度决策。但是，任何给定的批次要么仅包含其提示阶段的请求，要么仅包含token阶段的请求。提示阶段被认为更重要，因为它会影响 TTFT。因此，等待提示可以抢占token阶段。虽然这会导致更短的 TTFT，但它会大大增加 TBT 的尾部，从而增加 E2E 延迟。最后，还有混合批处理（图 ©）[23]。通过这种批处理，调度决策是在每次前向传递时做出的，并且提示和token阶段可以一起运行。这减少了对 TBT 的影响，但并没有消除它，因为与提示阶段一起调度的token阶段将经历更长的运行时间。

模型并行可用于将模型划分到多个 GPU 甚至多台机器上，以获得更高的效率和内存容量。LLM 推理通常使用流水线并行和张量并行。流水线并行 (PP) 将模型的各层划分到各个 GPU 上，同时将同一 GPU 上某一层内的所有算子和张量保留在同一个 GPU 上。张量并行 (TP) 将张量划分到各个 GPU 上，同时在每个 GPU 上复制所有层。流水线并行需要参与的 GPU 之间进行较低的通信，而张量并行需要每一层进行高带宽通信。通常，张量并行对于同一台机器内的 GPU 来说性能更好，并且通过高带宽互连（例如 NVLink [15]）。

随着 LLM 用例的兴起，一些云服务提供商扩展了基于 GPU 的产品，从而实现了大型 GPU 集群部署 [5][56][57]。这些 AI 集群中的每台机器通常由 8 个旗舰 NVIDIA GPU（A100 或 H100）组成。每个 GPU 通过高带宽 Mellanox InfiniBand 互连 [10][13] 连接到集群中的所有其他 GPU，形成高带宽数据平面网络。目前，云端提供的 InfiniBand 带宽为 25 - 50 GBps /每对 GPU [7][10]。

下表显示评估的模型。BLOOM [69] 和 Llama2 [71] 都是最先进的开源 LLM。这两个模型都是仅解码器、基于Transformer的模型。用每个模型中参数最多的版本，因为这些版本在产品级准确率方面最具代表性。实验中在一台配备 8 个 H100 [16] GPU 的机器上运行 vLLM [51] 的 BLOOM-176B 和 Llama-70B。

为了更好地理解这个思路，检查提示输入和生成的输出tokens数量分布。图 a 显示提示tokens数量的分布。由于编码 LLM 推理服务通常用于在用户编写代码时生成完成，因此其输入提示可以包含迄今为止编写的大量代码。因此，它的中位提示大小较大，为 1500 个 tokens。另一方面，对话服务具有更广泛的输入提示 tokens 范围，因为它取决于用户。此踪迹的提示token中位数为 1020 个tokens。图 b 显示了生成的 token 数量分布。由于编码服务通常仅在用户键入时在程序中生成接下来的几个单词，因此输出 tokens 的中位数为 13 个 tokens。另一方面，对话服务几乎具有双峰分布，中位数为 129 个 tokens。

其他还有一些因素考虑：批处理、延迟、吞吐量、内存和功率。

Splitwise是一种将 LLM 推理中的提示和生成阶段拆分到不同机器上的技术。如图显示了 Splitwise 的高级概览。

维护两个单独的机器池，用于提示和token处理。第三个机器池，即混合池，根据工作负载的需要进行扩展和收缩。所有机器都预先加载所选模型。当新的推理请求到达时，调度程序会将其分配给一对机器（即提示和token）。提示机器负责生成输入查询的第一个token，方法是处理提示阶段的所有输入提示tokens并生成 KV 缓存。提示机器还将 KV 缓存发送到token机器，token机器继续生成tokens，直到响应完成。在token机器上使用连续批处理来最大限度地利用它们。混合池中的机器使用混合连续批处理。

在较低的请求率下，目标是在 Splitwise 中实现更好的延迟，而在较高的请求率下，目标是避免由于提示和token机器池之间的碎片化而导致的任何性能或吞吐量降低。Splitwise 使用分层两级调度。集群级调度器 (CLS) （1） 负责机器池管理和路由传入的推理请求。机器级调度器 (MLS) （2） 维护待处理队列并管理每台机器上的请求批处理。

集群级调度器（CLS）维护提示、token和混合机器池 （3） 。Splitwise 最初根据预期的请求负载和输入/输出tokens 分布将机器分配给提示或token池。CLS 使用加入最短队列 (JSQ) 调度 [39],[85] 为每个请求分配一个提示机器和一个 token 机器。

机器级调度器（MLS）在每台机器上运行，负责跟踪 GPU 内存利用率、维护待处理队列 （4）、决定每次迭代的批次以及向 CLS 报告相关状态。

KV 缓存在请求的提示阶段生成，并且在token生成阶段不断增长。在 Splitwise 中，需要将 KV 缓存从提示机器传输到token机器 （5）以完成推理。此传输延迟是与 Splitwise 相关的主要开销。

如图是KV cache 传输示意图：

图 a 显示了在以序列化方式传输 KV 缓存时，单个请求批次的提示阶段、KV 缓存传输和token生成阶段的Gantt图。KV 缓存传输仅在提示阶段完成并生成第一个token后才开始。此外，它需要先完成，然后才能在token生成阶段生成下一个输出token。这直接影响推理的最大 TBT 和端到端延迟。传输所需的时间取决于 KV 缓存的大小（与提示tokens的数量成正比）以及提示机器和token 机器之间的互连带宽。即使使用快速的 InfiniBand 链路，大提示大小的传输开销也可能成为 TBT 的很大一部分。

在 Splitwise 中，将 KV 缓存传输与提示阶段的计算重叠来优化 KV 缓存传输。当 LLM 中的每一层在提示机器中计算时，也会生成与该层相对应的 KV 缓存。在每一层的末尾，触发该层的 KV 缓存异步传输，同时提示计算继续进行到下一层。图 b 显示了这种异步传输，它减少了传输开销。分层传输还可以实现其他优化，例如在 token 机器中更早启动 token 阶段，以及更早释放提示机器上的 KV 缓存内存。

分层 KV 缓存传输，与下一层的提示计算并行进行。这需要每层进行细粒度同步以确保正确性。因此，可能会产生性能干扰并增加 TTFT，尤其是对于较小的提示。但是，对于小提示，总的 KV 缓存大小很小，不需要分层传输来隐藏延迟。由于在计算开始时已经知道在一个批次中tokens数量，因此 Splitwise 选择了最佳的 KV 缓存传输技术。它对较小的提示使用序列化的 KV 缓存传输，对较大的提示使用分层传输。

为了在真实硬件上评估，在 vLLM [51] 之上实现 Splitwise 的 KV 缓存传输机制。实现开源 [1]。在 Microsoft Azure 上的两台 DGX-A100 和两台 DGX-H10 虚拟机 (VM) 上运行此修改后的 vLLM。这些机器与 InfiniBand 连接，DGX-H100 的带宽是其两倍（即 400 Gbps）。

Splitwise 技术的实现为机器分配了提示角色或token角色。当提示机器生成第一个token时，将 KV 缓存传输到token机器。用 MSCCL++ [11]（一种优化的 GPU 驱动通信库）来实现简单且逐层的 KV 缓存传输。

在实现中，提示机器使用 MSCCL++ 的零拷贝（zero-copy）单边 put 原语，在 KV 缓存数据准备就绪后，立即通过 InfiniBand 发送，而无需 token 机器发出任何接收指令。一旦为所有层发出 put，提示机器就会发出信号，token机器将等待该信号。借助信号完成的同步，使用与发送 KV 缓存数据相同的 InfiniBand 连接。在处理一批提示时，每个请求都会分配一个不同的信号，因为它可能会被路由到不同的token 机器。在 vLLM 中逐块发送 KV 缓存。为了最大限度地减少传输次数，还考虑 KV 块的连续性，只要它们使用相同的信号。

构建了一个模拟器来探索集群设计并大规模评估 Splitwise。模拟器代码开源 [20]。如图展示了模拟器的设计。

模拟器是事件驱动的，忠实地模拟了 Splitwise 机器池、调度程序、机器级内存和队列以及 KV 缓存传输。首先在具有各种输入/输出大小的目标硬件上对 LLM 进行分析 （1） 。基于配置文件，构建一个性能模型。模拟器将请求跟踪、SLO、性能模型以及集群和调度程序的配置作为输入 （2） 。对于评估，用提示和token大小分布。调整泊松到达率以增加和减少集群大小的负载（每秒请求数）。模拟器提供每个请求实现的指标（TTFT、TBT、E2E）和机器利用率水平 （3）。通过硬件实验交叉验证性能模型以确保准确性；还有超过 50K 次迭代的产品负载对模拟器进行了端到端验证，以确保保真度 （4） 。