TD-Pipe: Temporally-Disaggregated Pipeline Parallelism Architecture for High-Throughput LLM Inference

作者/机构: Hongbin Zhang, Taosheng Wei, Zhenyi Zheng, Jiangsu Du, Zhiguang Chen, Yutong Lu (Sun Yat-sen University)

A1 主要贡献

本文旨在解决流水线并行在以吞吐量为导向的大型语言模型(LLM)推理中面临的核心问题:由于prefill和decode阶段的工作负载不平衡以及复杂的数据依赖性,导致了大量的流水线气泡(pipeline bubbles),从而严重降低了性能。

核心问题:
现有流水线并行方法在LLM推理中效率低下。如图1所示,无论是将prefill和decode阶段分开批处理(Separate Batching)还是混合批处理(Hybrid Batching),都会因负载不平衡和数据依赖性产生严重的气泡。
* 负载不平衡:
1. 批间不平衡(Inter-batch imbalance): 不同批次的解码(decode)步骤因请求数量不同而工作量不同。
2. Prefill-Decode不平衡: prefill阶段的执行时间通常远长于单个decode步骤。

  • 数据依赖性:
    1. 层间数据依赖(Inter-layer data dependency): 数据必须在前一层处理完才能进入下一层。
    2. 解码步骤间数据依赖(Inter-decode-step data dependency): 下一个token的生成必须等待上一个token完成。
    3. 分块数据依赖(Inter-chunk data dependency): 在混合批处理中,单个prefill任务的切块之间引入了新的依赖。

即使采用分块prefill(chunked-prefill)和混合批处理的先进方法,也无法根本解决问题,GPU利用率依然很低(如图2左侧所示)。


图1: LLM流水线并行推理中的气泡。


图2: 流水线并行中的GPU利用率。

研究目标与创新点:
为了更好地利用流水线并行实现高吞吐量的LLM推理,本文提出了TD-Pipe,一个基于时间维度解耦的流水线并行架构(temporally-disaggregated pipeline parallelism architecture)

  • 核心思想: 在时间维度上将prefill和decode阶段解耦,即尽可能减少prefill和decode之间的切换频率。这样可以在各自的阶段内有效维持流水线各阶段的负载均衡,并大大解决数据依赖问题,从而显著提升GPU利用率(如图2右侧所示)。

  • 具体贡献:

    1. 识别了流水线并行在LLM生成任务中产生气泡的根本原因:工作负载不平衡和复杂的数据依赖。
    2. 设计了新颖的时间维度解耦流水线并行架构:该架构专门针对以吞吐量为导向的离线LLM推理,将prefill和decode阶段在时间上分离。
    3. 提出TD-Pipe系统及三套优化方法来解决新架构下的潜在问题:

      • 层次化控制器结构(Hierarchy-controller structure): 通过解耦调度与执行,更好地协调流水线中的设备。
      • 基于AI的贪婪prefill方法(AI-based greedy prefill): 通过预测输出长度和模拟内存使用,积极地执行更多prefill任务,延迟向decode阶段的切换。
      • 批间工作窃取方法(Inter-batch work stealing): 动态平衡不同批次在decode阶段的工作负载,减少气泡。
      • 时空强度对比方法(Spatial-temporal intensity comparison): 通过比较计算强度下降带来的性能损失和阶段切换气泡带来的性能损失,确定从decode切换回prefill的最佳时机。
    4. 对TD-Pipe进行了广泛评估,实验表明其在仅有PCIe互连的GPU节点上,相比现有的张量并行方法吞吐量提升高达1.91倍,相比现有的流水线并行方法提升高达2.73倍。

A3 背景知识

2.1 LLM生成任务

现代LLM生成任务遵循自回归解码过程,即根据输入序列预测下一个token。在此过程中,每个token位置都会生成称为KV缓存的中间状态。通过在内存中维护KV缓存,现代LLM推理系统可以消除这个迭代过程中的大部分冗余计算。根据KV缓存是否已生成,推理过程分为prefill和decode两个阶段。如图3所示,prefill阶段处理新输入的prompt(绿色部分)并初始化KV缓存(金色部分),通常包含许多token,并并行处理这些token以生成一个新token(蓝色部分)。接下来,得益于KV缓存,decode阶段的每一步只处理前一步生成的一个新token。LLM生成任务有几个关键特征:


图3: LLM生成任务的生成过程。

  • 首先,prefill和decode在计算上有显著差异。prefill阶段并行处理许多token,很容易成为计算密集型(compute-bound)。相比之下,decode阶段由于I/O量相似但每个请求只处理单个token,因此更多地受到内存带宽的限制。因此,prefill阶段只需很小的批量大小(batch size)就能饱和计算资源,而decode阶段则需要大得多的批量大小。
  • 其次,长度可变的特性带来了固有的不确定性。输入和输出的长度都在变化,且输出长度通常在推理完成前是未知的。

2.2 流水线并行在高吞吐量LLM推理中的机遇

2.2.1 高吞吐量LLM推理。尽管在线LLM服务备受关注,但在没有严格延迟约束的场景中,如RLHF的rollout阶段【【12,OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework,2024,arXiv,https://arxiv.org/abs/2405.11143】】和推荐系统的数据增强【【7 ,Data Augmentation using LLMs: Data Perspectives, Learning Paradigms and Challenges,2024,ACL,doi:10.18653/v1/2024. findings-acl.97】】,追求高吞吐量成为首要任务。

首先,实现高吞吐量LLM推理需要足够的内存容量。推理期间,模型权重和KV缓存都需要大量内存。常用的LLM参数量有7B、13B、30B、70B甚至更多,超出了大多数设备的内存容量,特别是像NVIDIA A10(24GB)、4090(24GB)和L20(48GB)这样的商用设备。除了权重,decode阶段还需要同时容纳数百个请求以饱和GPU资源,KV缓存是内存占用的主要部分。例如,Llama-30B【【25,Llama: Open and efficient foundation language models,2023,arXiv,https://arxiv.org/abs/2302.13971】】中单个token的KV缓存占 用1.52 MB,400个平均长度为300的请求大约需要178 GB内存。

其次,实现高吞吐量LLM推理需要在利用多个资源扩展内存容量时最小化协调开销。图4展示了部署LLM时广泛使用的流行多GPU架构,其中多个GPU连接到单个CPU。在这种架构下,通常有两种主要方法来提高内存容量:卸载(offloading)和并行化(parallelization)。


图4: 多GPU架构。


图5: 卸载和并行方法。


图6: 在多GPU架构中使用张量并行方法时,prefill阶段的执行时间分解。

2.2.2 卸载方法中的协调开销。如图5(a)所示,单个推理实例可以通过将数据卸载到二级存储来获得更多内存容量。通过在主机内存和GPU内存之间迭代交换数据,这使得decode阶段可以在给定时刻容纳更多请求进行处理。然而,即使在CPU与GPU比例为1:1的情况下,PCIe带宽也可能成为这种方法的瓶颈。在多GPU架构中,当多个GPU同时将数据卸载到主机内存时,会遇到更严重的PCIe争用问题。换言之,CPU和GPU之间巨大的协调开销使得这种方法不适用于高吞吐量LLM推理。

2.2.3 并行方法中的协调开销。通过将LLM分布到多个设备上,一个实例可以获得更多的内存容量。图5(b)和(c)展示了张量并行(TP)和流水线并行(PP)方法。TP将每一层分片到多个GPU上,模型权重和KV缓存均匀分布在GPU工作节点上。它需要频繁的通信,一个Transformer层就需要两轮all-reduce操作。我们在两个没有直连电缆的节点上对Llama-30B【【25,Llama: Open and efficient foundation language models,2023,arXiv,https://arxiv.org/abs/2302.13971】】进行了FP16的分布式推理案例研究。为了展示强大的可扩展性,我们减少模型中的层数以适应更少的设备,并报告了运行结果,如图6所示。值得注意的是,该模型由结构相同的层组成,因此减少层数不影响其计算或通信特性 。

在L20节点上,当设备数量从1增加到4时,总执行时间减少了1.84倍。尽管总执行时间减少了,但在4个设备上,通信时间占总时间的47.39%。在A100节点上,总执行时间仅减少了1.64倍,通信时间在4个设备上消耗高达53.9%。因此,张量并行方法严重依赖互连能力,并使计算资源长时间闲置。此外,某些高性能GPU通过专用电缆(如NVIDIA的NVLINK)实现直接通信,这显著提高了它们的通信效率,但同时也大大增加了节点价格。

与TP相比,PP按层切分模型,每个设备负责一部分层。它每隔几层只需要一次数据量小得多的点对点通信。尽管LLM推理中不平衡的流水线工作负载和复杂的数据依赖性仍然使其在延迟敏感场景中不是首选,但通信量的大幅减少使PP成为支持吞吐量导向的LLM推理的一个有前景的方法,特别是在没有高性能连接(如NVLINK)的商用硬件上。

2.3 现有的LLM推理流水线并行

现有的流水线并行主要在批处理方法上有所不同。为了充分利用现代GPU,批处理技术被广泛用于处理深度学习工作负载,其中多个样本同时处理以暴露高并行性并提供显著的硬件性能提升。LLM主要有分离批处理和混合批处理两种方式。

分离批处理(Separate batching)调度请求时,要么专门用于prefill阶段,要么专门用于decode阶段,不会在同一批次中混合它们。鉴于这两个阶段表现出不同的性能特征,即计算密集型和内存密集型,prefill阶段即使批量大小仅为1也能饱和GPU计算,而decode阶段则需要数百的批量大小。在任何时候,推理实例要么处于prefill阶段,要么处于decode阶段。相比之下,混合批处理(hybrid batching)将来自prefill和decode阶段的请求组合到同一个批次中,形成混合批次。

当与流水线并行集成时,这两种批处理方法各有优缺点,但都仍然存在严重的流水线气泡问题。使用分离批处理时,管理仅由prefill阶段请求组成的批次很简单,因为它们之间没有批间数据依赖,可以连续启动。同样,仅有decode请求的批次也便于实现批间负载均衡,因为每个decode步骤处理单个token,工作负载相对均匀。然而,在同一批次中混合prefill和decode请求会引入数据依赖和工作负载不平衡,从而导致严重的流水线气泡。而混合批处理通常与分块prefill方法【【1,Taming {Throughput-Latency} Tradeoff in {LLM} Inference with {Sarathi-Serve},2024,OSDI 24,】】结合使用,该方法将prefill拆分为多个块,从而能够实现更好的批间负载均衡。然而,1)由于decode步骤被添加到所有批次中,管理所有批次间的数据依赖变得至关重要;2)高度可变的输入和输出长度仍然导致严重的负载不平衡;3)此外,分块prefill引入了重复的KV缓存访问作为开销。

A2 方法细节

TD-Pipe的关键思想在于时间维度解耦的流水线并行架构,该架构在时间维度上将prefill和decode阶段分离开来。基本上,TD-Pipe旨在在多GPU架构中实现高吞吐量的离线LLM推理,特别是在没有GPU直连互连电缆的商用节点上。图7展示了TD-Pipe的整体设计。首先,TD-Pipe遵循一个层次化控制器结构,将系统抽象为一个中心化引擎和一个分布式运行时,分别负责控制和执行。通过将调度与执行解耦,它可以在流水线中异步协调多个设备。

接下来,TD-Pipe识别了时间维度解耦的流水线并行架构中的三个基本需求,并通过三种新颖的方法来解决它们。首先,基于AI的贪婪prefill方法通过预测输出长度和计算最佳点,积极地延迟并确定从prefill切换到decode的时机,从而在一个固定的内存容量内执行更多的prefill。其次,批间工作窃取方法减少了因不同批次中请求随机完成而引起的流水线气泡,并将工作负载从重载批次动态地重新分配到轻载批次。第三,时空强度比较方法通过比较计算强度降低带来的性能成本与切换引起的流水线气泡成本,来确定从decode切换到prefill的时机。最终,TD-Pipe实现了紧凑的流水线调度,在阶段切换时仅有极小的气泡。

3.1 时间维度解耦的流水线

如第2.3节所述,简单地将分离批处理集成到流水线并行中会导致严重的气泡问题,这主要是由于数据依赖和工作负载不平衡。为了解决这个问题,我们采用了一种时间解耦策略,在时间维度上分离prefill和decode阶段,从而最小化它们之间的相互干扰。在这种设计中,每个推理实例在很长一段时间内只处理一种类型的阶段,显著减少了流水线气泡。prefill阶段和decode阶段以生产者-消费者的关系运作,其中decode阶段依赖于prefill阶段的完成。因此,阶段切换仍然是必要的。为了进一步改进这个架构,我们的目标是减少阶段切换的频率,同时在每个阶段内保持高效率。

3.2 层次化控制器系统结构

为了让设备更好地协作,我们为TD-Pipe设计了层次化控制器系统结构。为了同时处理复杂的数据依赖并实现异步数据传输,如图7所示,它解耦了流水线任务的调度和执行,将中心化引擎抽象为控制平面,将分布式运行时抽象为执行平面。中心化引擎负责根据内存容量、请求状态和性能分析结果进行批处理调度,然后向分布式运行时的worker分发任务。分布式运行时负责具体的执行,每个worker只知道它应该与哪个设备通信以及通信什么数据。


图7: TD-Pipe的系统概览。


图8: 输出长度预测模型

3.2.1 中心化引擎。中心化引擎通过远程过程调用(RPC)管理和控制worker。它主要负责两个不同的阶段,即运行时初始化和执行调度。在初始化期间,它将子模型委托给worker,初始化模型的相关部分并将参数加载到内存中。执行调度处理请求的批处理,并向worker发送控制信息。对于prefill阶段,引擎打包prompt并直接发送给worker,直到触发切换条件。对于decode阶段,引擎将decode步骤打包成批次,并负责在每次迭代中删除已完成的请求。

3.2.2 分布式运行时。作为仅执行平面,分布式运行时遵循单程序多数据(SPMD)模型的思想,即每个worker只知道它应该计算什么数据,通信什么数据,以及应该与哪个设备通信。为了管理通信元信息,如world size和rank number,会初始化一个全局通信上下文。这样,在接收到中心化引擎分发的信息后,worker可以根据全局通信上下文知道自己的位置,并能异步地执行流水线阶段之间的点对点通信。

3.3 方法一:基于AI的贪婪Prefill

在执行prefill阶段后,TD-Pipe需要确定从prefill切换到decode的时机。为了减少prefill和decode阶段之间的切换频率,TD-Pipe倾向于在切换到decode阶段之前执行更多的prefill。然而,decode阶段也会生成KV缓存并消耗内存,盲目执行过多的prefill会导致在decode阶段频繁地进行重计算或卸载,从而导致性能显著下降。不幸的是,输出长度在请求完成之前是未知的,峰值内存使用量也是未知的。

为了确定在decode阶段应启动多少prefill而不大幅超出内存容量,我们提出了基于AI的贪婪prefill方法。如果我们能知道所有请求的输入和输出长度,我们就可以计算总内存使用量并确定切换时机。此外,在decode阶段,所有请求一起进行,一些请求生成新的KV缓存,而另一些请求则完成并释放相关的KV缓存。通过模拟整个过程中的内存使用情况,我们可以更积极地执行prefill。

为了准确预测每个请求的输出长度,我们遵循了µ-Serve【【21,Power-aware Deep Learning Model Serving with µ- Serve,2024,USENIX ATC 24,https://www.usenix.org/conference/atc24/ presentation/qiu】】中介绍的预测模型,如图8所示。该模型采用一个基于BERT的多类别分类器来估计每个用户请求的输出token长度。具体来说,它利用BERT最后一层[CLS] token的隐藏状态作为输入的聚合表示,然后通过一个两层前馈神经网络来预测输出长度的类别。这些类别被定义为范围,从[P0, P25)到[P99, +),表示从历史推理数据中得出的百分位区间。我们假设在给定场景中的推理输入具有很强的相似性,这使我们能够基于历史模式来训练模型。对于每个请求,预测的输出长度根据训练集中该类别的平均值来指定。值得注意的是,预测模型的执行开销极低,在整个推理过程中可以忽略不计。

有了输出长度预测,我们设计了一个动态规划(DP)算法,如算法1所示,以高效地模拟执行新prefill后后续解码步骤的内存使用情况。这个模拟在每个prefill完成后执行,不会引入显著的开销。为了简化计算复杂性,我们指定了未来的一些解码步骤用于决策,称为futurePoints。例如,我们只检查第32、64、96、...、992、1024个解码步骤的内存使用是否超过内存容量。在新增一个批次后,该算法通过检查预测的输出长度是否超过某些解码步骤来更新这些解码步骤(futurePoints)的KV缓存使用情况。有了未来的KV缓存使用情况,我们就可以通过检查是否存在一个futurePoint的内存使用大于内存容量来确定从prefill切换到decode的时机。

3.4 方法二:批间工作窃取

随着decode阶段的进行,不同批次中的请求会随机完成,导致不同批次的工作负载不断变化,从而产生不平衡的工作负载。由于decode步骤之间存在数据依赖,这可能导致流水线气泡。批间工作窃取方法负责确保在整个decode阶段,工作负载能够均匀地分布在各个批次之间。

首先,我们需要确定何时可以认为两个decode批次具有等效的工作负载。单个decode步骤的执行时间由两个因素决定:批次大小和KV缓存的长度。decode阶段通常处理大批次,数量常达数百,并且请求间的KV缓存长度也差异显著。这个巨大的空间使得精确比较两个批次间的执行时间变得困难。为了简化这个问题,我们认为批次大小是衡量批次间负载平衡的唯一指标。这是因为线性操作在大多数场景下通常主导Transformer模型的执行时间,其执行时间仅与批次大小相关。此外,大批次大小有助于平滑由序列长度变化引起的执行时间波动。


图9: 批间工作窃取方法。bs(wi)代表窃取后的批次大小,bs(wo)代表窃取前的批次大小。

在开始时,我们将请求分成与GPU数量相等的批次,每个批次包含相同数量的请求。随着decode阶段的进行,不同批次中的请求会随机完成,导致不同批次的工作负载不断变化,从而产生不平衡的工作负载。为了使它们平衡,批间工作窃取方法需要动态地在批次之间迁移工作负载。值得注意的是,批次调度器在某一时刻最多只能观察到一个批次的实际大小,而其他批次仍在执行中,一些请求可能已经完成。因此,该方法采用滑动窗口的思想来实时平衡批次间的工作负载。当一个批次完成并返回时,中央调度器首先检查该批次并移除已完成的请求。然后,通过从滑动窗口中维护的请求总数中减去已完成的请求数并计算平均值,得到一个平均批次大小。滑动窗口的长度等于GPU的数量,用于跟踪最近的批次大小。如果当前批次包含的请求数多于平均值,多余的请求将被暂时扣留,然后该批次被提交以进行进一步处理。在后续的迭代中,调度器将用扣留的请求补充其他批次。通过这个迭代过程,工作窃取方法逐渐重新分配工作负载,有效地缓解了负载不平衡。

图9展示了一个稍显极端的例子,说明了这种方法如何在一个4阶段流水线中逐渐实现负载均衡。最初,512个请求被均匀地分为四个批次,每批128个。第一次迭代后,批次0完成了48个请求,剩下80个。

如图9所示,通过对最后四个批次应用滑动窗口并减去已完成的请求数,我们得到平均批次大小为116。由于批次0中剩余的80个请求低于当前每批116个请求的平均值,所有80个请求都被重新提交。批次1中有8个请求完成,当前平均批次大小变为114。由于批次1中的请求数超过114,调度器从批次1中“窃取”6个请求,只提交114个。对于批次2和3,第一轮没有请求完成,调度器同样根据需要窃取请求,以使每个批次接近平均值。在下一次迭代中,这些被窃取的请求被重新分配给批次0,从而实现所有批次的工作负载均衡。在图9的下半部分,不平衡的工作负载迅速趋于平衡,减少的流水线气泡转化为性能增益。

3.5 方法三:时空强度对比

在时间维度解耦的流水线架构中,decode阶段扮演消费者的角色。随着decode阶段的进行,一些请求完成,导致批次大小减小,从而计算强度下降。如果计算强度持续下降,decode阶段的运行效率会很低。然而,如果我们在decode效率下降时立即切换到prefill阶段,频繁的切换可能会产生气泡并损害性能。为了处理这个两难问题,时空强度对比方法旨在确定从decode切换到prefill的最佳时机。

如图10所示,时空强度对比方法评估并比较继续执行decode阶段与切换到prefill阶段的计算效率。为了量化这两种效率,我们为保持在decode阶段定义了空间强度,为切换到prefill阶段定义了时间强度。


图10: 时空强度对比方法。

空间强度定义为当前decode性能与可达到的峰值性能之比。具体来说,我们使用足够大的批次大小来分析执行时间,并计算每个请求的平均执行时间。我们将高计算强度下每个请求平均执行时间的倒数称为峰值(Peak)。如图10左侧所示,对于每个批次大小,我们同样分析每个请求平均执行时间的倒数,我们将其表示为已实现值(Achieved)。空间强度则表示为:

$$Spatial\ Intensity = \frac{Achieved}{Peak}$$

然后,时间强度定义为1减去气泡比率。气泡比率是指如果现在发生切换,相对于下一个prefill阶段的持续时间,因气泡损失的时间比例。如图10右侧所示,对于decode阶段的每个时刻,我们计算潜在的气泡并模拟下一个周期的总时间。对于气泡,它是最长prefill与当前decode之间的差值。对于总时间,它是待处理的prefill、每个批次的一个decode步骤(甚至可以忽略)以及气泡的累积。由于待处理的prefill是确定的,我们可以轻松计算它们的执行时间。因此,时间强度表示为:

$$ \text{Temporal Intensity} = 1 - \frac{bubble}{total} $$

基于这两种强度,我们设计了决策策略如下:当空间强度小于时间强度时,系统切换到prefill阶段。否则,它将继续执行decode阶段。

A4 实验环境

实现: TD-Pipe基于vLLM 0.5.3【【14,Efficient memory management for large language model serving with pagedattention,2023,SOSP,】】实现。

硬件配置:
* 节点1: 4块NVIDIA L20 GPU(48GB显存),通过PCIe交换机连接,峰值All-Reduce带宽为14.65 GB/s。
* 节点2: 4块NVIDIA A100 GPU(80GB显存),通过PCIe交换机连接,峰值All-Reduce带宽为14.82 GB/s。
具体配置见表1。

表1: GPU配置

软件配置:
* PyTorch v2.3.1
* vLLM v0.5.3
* CUDA 12.4
* NCCL 2.20.5

模型架构:
* Llama2-13B-chat
* Qwen2.5-32B-Instruct
* Llama2-70B-chat
具体参数见表2。其中32B和70B模型使用了分组查询注意力(GQA)技术。

表2: 模型规格

数据集与用途:
* 数据集: 使用流行的对话数据集shareGPT V3【【9,shareGPT v3,2024,Hugging Face,https://huggingface.co/datasets/ anon8231489123/ShareGPT_Vicuna_unfilteredr】】。
* 数据处理: 过滤输入长度小于1024 token的句子,生成了86,612个输入输出对。
* 用途:
* 输出长度预测模型训练: 60%样本用于训练,20%用于验证。
* 性能评估: 随机抽取5,000个输入句子作为请求进行测试,每次测试运行约半小时。

  • 评估指标: 吞吐量(tokens/s)。

A5 实验结果

4.2 总体性能评估

实验内容:
在L20+13B、L20+32B、A100+32B和A100+70B四种配置下,将TD-Pipe与四种基线方法(TP+SB, TP+HB, PP+SB, PP+HB)进行比较,GPU数量从1扩展到4。

实验结果:
如图11所示,TD-Pipe在几乎所有情况下都优于其他方法,尤其是在4个设备时。
* 在所有4-GPU配置中,TD-Pipe的吞吐量比TP+SB最高提升1.91倍,比TP+HB最高提升1.90倍,比PP+SB最高提升2.73倍,比PP+HB最高提升2.21倍。
* TD-Pipe展示了超线性的加速比。例如,在L20+32B配置下,GPU从2个扩展到4个时,吞吐量增长了2.97倍。


图11: 整体性能。

分析结论:
* TD-Pipe的超线性加速得益于计算能力和内存容量的双重提升。增加的内存容量可以容纳更多的KV缓存,从而显著提高decode阶段的计算强度。
* 其他流水线并行方法(PP+SB, PP+HB)受内存容量影响较小,且更长的流水线阶段加剧了它们的气泡问题,因此性能提升有限。
* TD-Pipe在A100节点上的性能优势比L20节点上更明显(相比TP方法)。这是因为A100的计算和内存带宽更强,使得TP方法更受限于互连带宽,而TD-Pipe对互连的依赖性小得多。同时,A100更大的内存容量也让TD-Pipe受益更多。

4.3 KV缓存内存使用情况

实验内容:
在4*A100+70B配置下,追踪推理过程中实际KV缓存使用量与总分配KV缓存内存的比例。

实验结果:
如图12所示,内存使用率先在prefill阶段(蓝色)持续增长至接近饱和,然后系统在prefill和decode阶段(黄色)之间交替。在decode阶段,内存使用量先增长,然后随着请求完成而下降。


图12: KV缓存内存使用波动情况。

分析结论:
内存仅在很短的时间内被完全占用,这表明基于AI的贪婪prefill方法能有效管理内存,避免了长时间的内存饱和。

4.4 消融实验

4.4.1 方法一:基于AI的贪婪Prefill
* 实验1: Prefill到Decode的切换时机
* 内容: 将AI方法与一个固定的“KV缓存占用率”超参数进行比较,当占用率达到阈值时切换。
* 结果: 如图13所示,TD-Pipe的AI方法优于所有手动选择的切换点。
* 结论: AI方法能动态适应变化的请求和硬件配置,找到更好的切换点,无需人工调优。


图13: Prefill到Decode切换的消融研究。

  • 实验2: 输出长度预测器性能
    • 内容: 评估预测器的准确性和开销。
    • 结果: 单个请求的预测准确率在0.52到0.58之间。更重要的是累积长度预测误差,如图14所示,随着请求数量增加,误差迅速减小。例如,预测256个请求的总输出长度时,误差仅为2.84%-6.18%。预测器开销极小,占总处理时间的不到0.16%。
    • 结论: 累积预测精度足够高,且开销可以忽略不计。


图14: 不同请求数量下的累积误差。

4.4.2 方法二:批间工作窃取

  • 实验内容: 比较启用(wi)和禁用(wo)批间工作窃取方法时的吞吐量。
  • 结果: 如图15所示,在L20+32B和A100+70B配置下,启用该方法分别带来了1.14倍和1.07倍的吞吐量提升。
  • 结论: 验证了在decode阶段动态平衡负载的有效性。


图15: Decode批次负载均衡的消融研究。

4.4.3 方法三:时空强度对比
* 实验内容: 将时空强度对比方法与一个固定的“请求完成率”超参数进行比较,当完成率达到阈值时从decode切换回prefill。
* 结果: 如图16所示,时空强度对比方法始终能达到最高的吞吐量。
* 结论: 验证了该方法在确定decode到prefill切换时机上的有效性。


图16: Decode到Prefill切换的消融研究。

A6 结论

为了在带宽受限的商用硬件上实现高吞吐量的LLM推理,本文提出了时间维度解耦的流水线并行架构。围绕这一新颖架构,我们识别了关键问题并引入了LLM推理系统TD-Pipe。TD-Pipe的核心是将prefill和decode阶段在时间维度上解耦。为了更好地利用该架构,TD-Pipe将执行平面和控制平面解耦,以实现更好的设备协调。此外,TD-Pipe设计了优化prefill-to-decode和decode-to-prefill切换时机选择的解决方案,并在decode阶段动态平衡批次间的工作负载。实验结果表明,在PCIe节点上扩展到4个GPU时,TD-Pipe的性能分别比现有的张量并行和流水线并行方法高出1.91倍和2.73倍。