无干扰推理：混合下游工作负载的 分解LLM 推理

24年1月来自中科院和华为云的论文“Inference without Interference: Disaggregate LLM Inference for Mixed Downstream Workloads”。

基于 Transformer 的大语言模型 (LLM) 推理服务现已成为许多云服务的骨干。LLM 推理包括预填充阶段和解码阶段。然而，现有的 LLM 部署实践往往忽视了这些阶段的不同特点，从而导致严重的干扰。为了减轻干扰，根据推理请求的特点仔细调度和分组推理请求。通过三大支柱，在 TetriInfer 中实现了这个想法。首先，它将提示分成固定大小块，加速器总是接近其计算饱和极限。其次，它分解预填充和解码实例，以便每个实例都可以独立运行。最后，它用一个智能两级调度算法并增强预测资源使用情况，避免解码调度热点。结果表明，TetriInfer 在性能/美元方面大幅提高了第一个token时间 (TTFT)、作业完成时间 (JCT) 和推理效率，例如，它使用的资源减少了 38%，同时平均 TTFT 和平均 JCT分别降低了 97% 和 47% 。

自 ChatGPT 以来，基于大语言模型 (LLM) 的服务已在日常生活中发挥着至关重要的作用。要运行推理请求，LLM 模型将首先获取用户输入以生成第一个 token（称为预填充阶段），然后以自回归方式逐个 token 生成输出（称为解码阶段）。

有多种方式可以与 LLM 交互，从简单的聊天到更复杂的下游任务，例如文档摘要、内容创建等。因此，支持 LLM 的服务可处理具有截然不同属性的推理请求，这些属性可分为两个维度：预填充阶段的输入提示长度和解码阶段的生成 token 长度。如图所示，摘要任务具有较长的输入提示和较短的生成tokens，而上下文创建任务则相反。不同下游任务的tokens长度可能相差两个数量级以上。

鉴于来自各种下游任务的 LLM 推理请求存在显著差异，第一个研究的问题是，这些推理请求在一起运行时的表现如何？为了回答这个问题，大量测试混合了不同长度的 LLM 预填充和解码请求。不幸的是，所有组合都存在严重干扰。例如，混合预填充请求可能导致速度减慢 10 倍，组合预填充和解码请求可能导致速度减慢 5 倍，而混合不同长度的解码请求可能会导致 16% 的吞吐量损失。避免干扰的一个简单解决方案是静态地为每个下游任务配置资源。

鉴于 LLM 服务基础设施的高成本，这种解决方案是不切实际的。为此，第二个研究的问题是，如何构建一个最小化干扰的分布式 LLM 推理服务系统？

LLM 推理，是根据输入提示生成输出tokens序列的过程。此过程包含两个阶段：预填充和解码。预填充阶段输出第一个token并生成K-V缓存以供将来解码 [21]。解码阶段使用之前的 KV 缓存以自回归的方式逐步生成新tokens。通常，预填充阶段受计算限制，而解码阶段受内存限制 [33]。如图所示，实验结果表明，一旦加速器在一定数量的tokens上达到饱和（称为加速器饱和阈值），预填充阶段的吞吐量就会保持平稳。解码阶段的吞吐量会随着批次大小的增加而继续增加，但一旦内存带宽饱和就会达到稳定水平。

为什么存在干扰？根本问题在于，当前的 LLM 部署实践没有考虑到 LLM 预填充和解码阶段所表现出的不同特征。具体来说，预填充阶段类似于计算繁重的批处理作业，其计算量与输入提示长度成二次方关系。解码阶段类似于内存密集型、延迟-关键型任务，其资源使用量与生成的tokens长度成亚线性关系 [33]。

在测试中观察到的干扰，是经典的系统问题。运行预填充请求，会导致严重的减速，因为这是在饱和的硬件继续添加计算繁重的作业。结合预填充和解码请求会对两者造成损害，因为系统在同时运行批处理和延迟关键型作业。混合解码请求会导致吞吐量下降，因为不知道内存带宽和容量使用情况，从而导致争用和队头（head-of-line）阻塞。

TetriInfer 是一个旨在对抗干扰的 LLM 推理服务系统。如图展示其工作流和架构图：(a) 比较现有系统和 TetriInfer 的执行时间线，现有系统有两个节点运行混合预填充和解码，TetriInfer 将预填充和解码实例分开，这允许在不同的按需解码实例之间对长解码任务进行负载平衡；每个时间线包含四轮（R1 到 R4），预填充和解码框的长度表示它们的序列长度，解码框的宽度表示其资源使用情况；更宽的解码框表示存在较长的生成tokens，从而导致更大的资源使用率和解码延迟；(b) 显示 TetriInfer 的架构，其中突出显示了四个核心模块：集中控制平面、预填充实例、解码实例和长度预测模型。

集中控制平面。它由一个全局调度程序和一个集群监视器组成。全局调度程序根据负载将请求发送到预填充实例，并接收来自解码实例的流输出。集群监视器从预填充和解码实例收集统计信息，并定期将负载信息广播到预填充实例。它添加、删除和翻转预填充或解码实例。

预填充实例。它们仅运行 LLM 推理请求的预填充阶段。每个预填充实例都有一个本地调度程序、一个长度预测器、主 LLM 引擎和一个调度程序。所有请求都经过四个步骤。1）本地预填充调度程序根据预定义的策略对请求进行排序。2）长度预测器运行预测模型来推测请求的解码长度，然后使用它们来估计解码阶段的资源使用情况。3）主 LLM 引擎将所有请求划分为固定块。4）对于每个请求，调度程序运行解码间负载平衡算法来选择一个解码实例，然后将生成的 KV 缓存转发给它。

解码实例。它们实际上与预填充实例分离，并且仅运行 LLM 推理请求的解码阶段。每个解码实例都可以接收来自任何预填充实例的请求。它运行一个本地调度程序，该调度程序具有三个预定义策略，用于选择要在主 LLM 引擎中运行的解码请求。

长度预测模型。预测模型是一个小型 LLM 模型，经过离线微调，用于预测 LLM 推理请求的生成长度。TetriInfer 的预填充调度程序和解码实例的本地调度程序利用推测的信息来调度解码实例并避免测量热点。预测模型很小，部署在所有预填充实例上。

预填充实例的调度程序对于改善预填充阶段的延迟和吞吐量至关重要。调度程序维护一个原始请求队列，该队列存储来自全局调度程序的请求，以及一个存储排序请求的调度队列。在这项工作中，实现了三种调度程序策略：先到先服务 (FCFS)、最短作业优先 (SJF) 和最长作业优先 (LJF)。

如图是预填充调度程序策略。左侧显示四个原始推理请求（R1 至 R4）。右侧显示使用 FCFS、SJF 和 LJF 的调度请求。这种分块版说明了切片和合并（C1 至 C4）。具体而言，FCFS 保持原始请求到达顺序。提示tokens按顺序划分并合并为块。此策略最容易实现，并且最适合具有相似提示长度的推理请求。但是，当请求具有较长的提示时，FCFS 可能导致队头阻塞和较高的平均作业完成时间 (JCT)。这是有问题的，因为 LLM 推理请求之间的长度差异超过三个数量级。

如图是预填充预测模型的微调和预测流程。目标模型是想要预测其解码行为的模型。预测模型是训练的模型，不涉及在线微调。在此过程中，训练预测模型（以红色表示）以推测特定目标模型（以蓝色表示）的解码行为。预测模型的微调涉及三个关键步骤。首先，组装一个从公共数据集继承的仅提示训练数据集、一个大型目标 LLM 模型（例如 OPT-13B）和一个用于预测模型的分类模型（例如 125M OPTForSequenceClassification [16]）。其次，向目标 LLM 模型发送训练提示，该模型生成响应。随后，将生成的响应分类到具有选定粒度的固定大小存储桶中。例如，使用粒度 100，token长度在 0 到 200 之间的响应 tokens 类别为 0，在 200 到 400 之间的响应 tokens 类别为 1，依此类推。这些标签与训练提示配对以创建新的数据集。最后，将新数据集划分为训练部分和评估部分，然后继续使用该数据集训练和评估预测模型。

一旦主 LLM 完成其预填充阶段，预填充的 KV 缓存就会在加速器的内存（例如 GPU 的 HBM）中生成。目标是将此缓存传输到所选解码实例的加速器内存，不管系统部署在哪个硬件平台上。这很有挑战性，因为预填充和解码实例之间存在多个物理数据链路，每个链路都需要不同的软件堆栈。

将现有的物理数据链路分为三类，如图所示。第一种称为Direct，其中加速器具有直接连接的高速链路，例如 NVLink [40] 或 HCCS [13]。可以使用低层内存复制原语 [14, 26] 或收集库 [28] 通过这些链路传输数据。第二种称为 Direct-NIC，其中加速器通过其配套 NIC 进行通信。可以使用定制库 [27] 通过 PCIe 和以太网（或 Infiniband）传输数据。第三种称为Indirect，没有直接链接，加速器必须通过其配套的 CPU DRAM 反弹数据，从而产生额外的内存复制。图中还将利用上述数据链接的网络堆栈分为单侧和双侧，类似于 RDMA 的分类。GPU 或 NPU 等加速器可以进行单侧内存访问，因为它们具有低级原语，例如设备之间的直接内存复制 [14, 26]。为了导航复杂的物理数据链接并确保 TetriInfer 在部署后始终能够使用性能最高的链接，设计一个统一的网络传输抽象来利用图中列出的不同网络堆栈选项。堆栈公开发送、接收、读取、写入等 API。调度程序调用这些 API 将 KV 缓存传输到远程解码实例。

在解码实例中，一旦选择了实例，将执行图中所示的步骤。要翻转预填充实例，全局调度程序会停止转发请求，然后向其发送翻转请求。预填充实例将等到所有排队的请求都耗尽。然后，翻转实例。翻转解码实例稍微复杂一些。全局调度程序通知所有预填充实例停止将请求转发到选定的解码实例，并通知解码实例完成翻转。请注意，实际翻转快速而简单。它涉及更改内部变量，而无需重新启动流程或重新加载模型。在当前的实现中，两个实例翻转大约需要 5 到 7 毫秒，不包括动态耗尽时间。

从头开始用 Python 实现 TetriInfer 的集中控制平面。采用基于 vLLM [21] 的预填充和解码实例。除了统一网络堆栈之外，大多数核心模块都是用 Python 实现的，统一网络堆栈使用 C++ 与低级 API 和 IB Verbs 交互以进行网络传输。此外，还实现了一种基于共享内存的通信机制，可以跨 Python 和 C++ 语言快速传输命令。微调部分使用 HuggingFace Transformer [16] 提供的 Trainer API。

预填充或解码实例是一个可部署的单元，部署时由两个进程组成。对于预填充，它有一个运行调度程序、长度预测器和主 LLM 的 Python 进程，以及一个运行调度程序和网络堆栈的 C++ 进程。对于解码，它有一个用于运行调度程序和主 LLM 的 Python 进程，以及一个运行网络堆栈的 C++ 进程。

由于高端硬件可用性有限，当前的实现仅支持使用sockets的Indirect类型。为了评估 TetriInfer 在不同硬件配置下的性能，实现了一种模拟机制来模拟不同的网络带宽。该机制的工作原理如下：对于给定的一组请求，首先离线运行它们的预填充阶段以获取其预填充的 KV 缓存。在测试之前，将这些预填充的 KV 缓存加载到解码实例的本地内存中。当测试开始时，预填充实例仅将请求元数据传输到解码实例，而不传输实际的预填充 KV 缓存。随后，解码实例计算 KV 缓存传输的延迟并进行相应等待。此延迟根据特定的模型架构和要模拟的硬件带宽计算。