大家好，最近Kimi开源了它的推理架构Mooncake的技术报告，让分离式推理架构的关注度一下升了起来。所以在这个系列中，我打算写一写关于分离式推理架构的一些有趣的优化知识。对于这个架构，我之前也只是了解一些，并没有深入探究过，所以在这个系列中我也和大家一起学习，一起发现好玩的东西。

本篇作为该系列的第一篇， 选择DistServe这个分离式架构入手，选择它的原因是 ：

• 这篇论文中通过各种实验和数学建模，很好论述了“为什么要用分离式架构”这一点。很适合初次了解这个架构的朋友阅读。这也是本文强调的重点

• 这篇论文代码是开源的，本文在写作时，也借鉴了开源代码的一些内容

• 调度策略比较简单（FCFS），也没有做抢占之类的操作（所以本文的重点也不会放在这里）

【全文目录如下】
一、LLM推理的两阶段及评估指标
二、为什么需要prefill和decode分离
三、分离架构下prefill和decode的优化方向
四、实操：代码实践中的prefill和decode并行策略

(每个大标题下还有很多细节，大家可以看正文，这里偷个懒，就不把小标题打出来了)

【如有帮助，欢迎大家点赞收藏和在看～～～～～】

一、LLM推理的两阶段及评价指标

1.1 LLM推理的两阶段

• prefill阶段 ：把整段prompt喂给模型做forward计算。 prefill阶段结束后，模型产出第一个token （例如图中t3）

• decode阶段 ： 一个token一个token地产出response 。（例如图中逐一产出t4，t5，t6）

既然LLM的推理由这两个阶段组成，那么一个LLM系统的性能评估同样也和这两个阶段密切相关。解析来，我们来看两个分别用于评价这两阶段性能的指标。

1.2 prefill性能评估指标：TTFT

TTFT（Time To First Token），表示生成第1个token所用的时间，这是prefill阶段的重要评估指标之一。

在实践中，我们一般会人为设置一个 TTFT SLO（Service Level Objective） ，这个指标是我们给系统提的要求，只有系统达到了这个指标，我们才认为它的性能达标了。举例来说， 假设我们定P90 TTFT SLO = 0.4s ，那么这意味着我们对该系统的要求是：“90%的request的TTFT值都必须<=0.4”

1.3 decode性能评估指标：TPOT

TPOT（Time Per Output Token），产出每一个response token所用的时间，这是decode阶段的重要评估指标之一。

同样，我们也有一个 TPOT SLO指标 ，这是人为给系统定下的性能要求。

二、为什么需要做PD分离

在快速回顾LLM推理的两阶段及它们的评估指标后， 现在我们来看一个重要的问题：为什么需要做PD分离？

2.1 分离架构速览

在一些推理框架中，prefill和decode是合并在一起的，你可以通俗理解成一块gpu上我既做prefill又做decode。我们以vllm框架为例 ：在一张卡上它的prefill和decode是交替进行的，即vllm的调度器会根据存储使用情况、请求池中的请求状态等要素，决定在模型推理的每1个阶段是做prefill还是做decode。

那 分离式的框架长什么样呢？我们直接来看DistServe提供的一张架构简图：

从这张架构简图中， 我们不难发现，在分离式框架中prefill和decode不再共享一块/一群gpu，而是你做你的prefill，我做我的decode 。prefill实例计算完毕后，把产出的KV Cache传送到decode实例，然后decode实例继续做推理。这里我们再明确一些术语：

• Prefill instance ：prefill实例。 一个prefill实例中包含一个完整的模型 ，它可能占据1或多块gpu。

• Decode instance ：decode实例。定义同上，只是它与prefill实例不再共享gpu。

介绍到这里，你肯定开始疑惑了： 从直觉上看，分离式架构相比于合并式架构，多加载了模型副本（耗显存），同时还涉及到gpu间的KV Cache传输（耗时间），怎么看都好像比不上合并式架构，那么我们做分离的原因何在呢？

接下来，我们就来回答这个问题。

2.2 TTFT和TPOT间的balance

假设你现在有一个requests pool，里面装了状态不同的requests，有正在等待做prefill的，也有刚做完prefill等待做decode的，而你面前只有1张GPU。由于prefill和decode在做fwd阶段的差异性，如果不考虑特殊的优化方法，在模型每1个推理阶段，你要么全做prefill，要么全做decode。 现在你需要做一个决策：在下1个推理时刻，我该做prefill，还是做decode？

你观察了一下requests pool，发现里面塞满了等待做prefill的请求，你心想“要不还是赶紧多做些prefill，好歹给用户返回第一个token，防止他们生气了直接不用我的产品了”。所以你在接下来的时刻中一直在做prefill， 由于gpu资源是有限的，你做了prefill就无法同时做decode了，你其实是在牺牲TPOT保全TTFT，反之亦然。

所以这时，你产生了一个疑问：当prefill和decode合并在一起的时候，似乎无论调度策略怎么调整，TTFT和TPOT这两个指标都存在强耦合关系。因此如果我把这两个阶段分开，让TTFT和TPOT各自独立发展优化，会不会得到一个更好的结果？

DistServe做了一个实验来验证你的猜想，接下来我们就来看这个实验结果。

2.3 分离VS合并：一个实验结果

实验配置：

• 1张A100 80G的gpu
• 13B LLM，input_length = 512, output_length = 64

SLO（人为定义的系统性能达标要求）：

• P90 TTFT = 0.4s，人为定义在prefill阶段90%的请求的TTFT必须达到0.4s（上面那张图黑色虚线部分）
• P90 TPOT = 0.04s，人为定义在decode阶段90%的请求的TPOT必须达到0.04s（下面那张图的黑色需要部分）

我们先只看图中的蓝线部分：

• 在蓝线部分对应的实验中，我们采用PD合并的架构
• 横轴表示每秒到达这块gpu上的请求数量，记为rps(requests per second)
• 随着rate增加，我们统计每个请求的TTFT和TPOT指标，进而绘制出蓝线。
• 将蓝线和我们设定的SLO（黑虚线）对比发现，如果我们同时达到设定好的TTFT SLO和TPOT SLO，我们的最大rps = 1.6，我们也记这个指标为goodput。

再来看图中的黄线和绿线部分：

• 还是这块gpu，现在我们做两次实验，一次实验只让它处理需要做prefill的请求（黄线），一次实验只让它处理需要做decode的请求（绿线）
• 按照前文介绍的流程，我们得到在这块gpu上只做prefill时，它的goodput=5.6；只做decode时，它的goodput=10。

从这个实验中，我们可以发现，一张卡只做prefill或者只做decode的goodput，都要比它既做prefill又做decode的goodput要高，这也验证了我们前文的猜想：合并式架构可能会同时损害TTFT和TPOT。

假设现在我们不考虑KV cache传输的网络成本，观察上述实验结果可发现：纯decode的goodput是纯prefill的goodput的2倍，所以我们可以估算：如果我拿2张卡纯做prefill，1张卡纯做decode，那么我这三卡构成的架构的goodput就为10，也即单卡goodput = 3.3。这是合并式goodput（值为1.6）的2.1倍！

再换一个角度粗糙地看：分离式架构中，我们三张卡承受起了10 reqs/s的流量；而合并式架构中由于单卡流量承受为1.6 reqs/s，所以我们需要6张卡才能承受起10 reqs/s的流量。 因此回到2.1结尾中我们的疑问，乍一看分离式架构似乎更耗卡了，但实际中如果优化得当，它其实是比合并式架构更加省钱的。

三、分离架构下prefill/decode的优化方向

通过上面的分析，我们具象化体会到prefill和decode分离的好处。 在本节中，我们再进一步，详细来看当我们确定可以使用分离架构后，prefill和decode阶段都有哪些优化方向。

3.1 算力与存储的独立优化

• prefill阶段：拥有计算受限的性质（compute-bound） ，特别是在请求流量较大，用户的prompt也比较长的情况下。prefill阶段算完KV cache并发给deocde阶段后，理论上prefill就不再需要这个KV cache了（当然你也可以采用LRU等策略对KV cache的保存做管理，而不是一股脑地清除）。

• decode阶段：拥有存储受限的性质（memory-bound） ，因为token by token的生成方式，decode阶段要频繁从存储中读取KV Cache，同时也意味着它需要尽可能保存KV cache。

对comput-bound和memory-bound不理解的朋友，可以参见我之前的 这篇文章 。

正是由于这两个阶段这两种特性，我们在为其做硬件分配时，也有了更加灵活的空间，而不是只使用一套型号完全一致的硬件。 也就是在分离式框架下，计算和存储可以朝着两个独立的方向做优化。

3.2 batching策略的独立优化

我们依然先展示DistServe所作的一个实验结果。

这个实验比较了在不同batch size下，prefill和decode阶段的吞吐量情况（tokens/s，即每秒能处理的tokens数量）不同颜色的线条代表不同大小的input_length。

• prefill阶段：随着batch size的增加，吞吐量的增长趋势却趋于平缓 。这正是因为prefill阶段是compute-bound的，当batch中的总tokens数大于一个阈值的时候，prefill阶段产生compute-bound，使得吞吐量的增长趋势变缓（这里我们不考虑显存问题，因为在该实验中硬件显存足够这些不同batchsize和input_length的数据使用）。

• decode阶段：随着batch size的增加，吞吐量的增长趋势越来越显著 。这是因为decode阶段是memory-bound的，即相比于计算，读写数据的时间要更多。所以在decode阶段中，如果我们能提升batch size，就能把计算强度提起来，吞吐量就上去了。

通过这个实验，我们不难得出结论：prefill和decode不适合用一样的batching策略，所以当两者分离后，我们可以在batching上继续做优化。

3.3 并行策略的优化

在合并式框架中，prefill和decode共享一个模型，这意味着它们共享一样的dp/tp/pp/....等并行策略。考虑2.3中的实验结果，我们根据实验指标，认为可以设计一个 2个prefill实例（dp=2），1个decode实例（dp=1）的分离 系统。所以我们自然而然想到，如果prefill和decode分离了，那是不是意味着它们的并行策略可以各不相同呢？

由于在前面的定义中，我们认为一个prefill/decode实例包含一个完整的模型副本。所以在本节的并行策略中，我们只考虑tp + pp（即1个实例只通过tp +pp划分） 。这样一来，只要对这两个阶段，分别确认好最佳的tp + pp策略，我们就可以用这个策略复制多个prefill/decode实例了（dp维度）

明确了这一点，接下来，我们通过一个有趣的实验 + 数学建模的方式，来看一下prefill和decode在并行策略（tp, pp）上的倾向性。

（1）prefill阶段的并行策略

怎么样可以知道不同并行策略对prefill阶段的影响呢？最简单的方法当然是通过做实验。实验我们肯定是会做的。 但在展示实验结果之前，我们尝试一个有趣的思路：能否拟合一个数学模型，来评估不同并行策略对prefill阶段的TTFT指标的影响？

如果我们能做这样的数学建模，同时用实际实验结果去论证数学模型的有效性， 那么之后我们的系统在选择最佳并行策略时，就不需要通过做实验（比如模拟假数据，送入模型实际做fwd）的方式去搜索最佳的并行策略了。我们完全可以用自己的数学模型去估算一个结果，这样岂不是能降低系统在初始化时的耗时？

在DistServe中，在研究不同并行策略对prefill阶段TTFT的影响时，就使用了“排队论”建立了这个数学模型。不过由于这个模型中一些强假设性（我们马上会讲这点），它并没有用在最终的代码实践中（代码实践换了一种建模方式，我们马上也会谈）。不过因为我觉得这个基于排队论的建模挺有意思的，所以我在这节中也会详细来讲。

在下面的讲解中，对于prefill阶段，我们分别对于单卡，2卡pp，2卡tp的情况做基于排队论的建模 （没错，这里我们先不谈tp + pp 的混合情况，这里我们主要想看对于prefill阶段，它对纯tp和纯pp的倾向性）。

单卡下prefill TTFT

我们假设队列中的请求服从排队论中的M/D/1模式，这个模式的含义如下：

• M：请求的到达过程遵从柏松分布 。即请求的到达系统是相互独立的（不存在互相干扰）且单位时间内各请求到达系统的概率是等可能的。

• D：假设所有请求做prefill的处理时间相同 。这个处理时间是指请求被喂给模型一直到模型吐出结果的时间。这就是我们前文说的“强假设”情况。

• 1：可以理解成只有1张gpu

在这个模式下，单卡下prefill阶段的avg_TTFT可以被建模成：

• D：即前文所说的单个请求处理时间
• R：即遵从柏松分布的请求到达率（每s有多少个请求到达gpu）上。

我知道你现在一定对公式（1）一头雾水，不要急，我们这就来详细解释它。想象一下，对于队列中的某个请求，它的TTFT要怎么计算呢？稍一动脑，我们不难从直觉上给出：

一个请求的TTFT = 这个请求在队列中等待的时间 + 单个请求被处理的时间 = 排在这个请求之前的请求数量 * 单个请求被处理的时间 + 单个请求被处理的时间

这里，单个请求被处理的时间 = D，而根据排队论，我们可以从概率论的角度推算出（推导过程略）排在这个请求之前的请求数量 = 。这样，我们就得到了公式（1）

最后注意，在M/D/1理论的假设下，RD一定是小于1的，如果RD大于1，则意味着这个系统的队列长度会趋向无限（即处理速度比进入速度慢，且进入是不停歇的），那么这时候就不再适合用这个排队理论了。

2-way PP下prefill TTFT

现在，我们假设有两张卡，我们对模型做2-way pp，则此时prefill阶段的TTFT建模如下：

下标inter表示层间切分，在这篇文章里我们都表示成pp，其余指标基本不变，只是D变了，具体来说有： 。哦？这是什么意思呢(这可能也是很多朋友读文章时的困惑)

• ： 这个Ds表示这个请求的处理时间 ，也还是为这个请求被喂进模型到这个模型吐出相应结果的时间。所以它和D是近似的。

• ： 这个Dm就不太相同了 。想象一下排在这个队列中最前面的两个请求。在单卡情况下，当第一个请求进入模型后，第2个请求需要等待D的时间才能被处理。但是在2-way PP的情况下，第1个请求过完第1张卡，第2个请求就可以跟上了，所以第二个请求理论上的等待时间预估为 ，这也就是我们的Dm。 再严谨一点，Dm可以认为是pp过程中最慢的那个stage的处理耗时。

2-way tp下prefill TTFT

同样假设有2张卡，做2-way tp，则我们有：

这里的K可理解成是一个折损因子 。理论上，如果不考虑因为tp产生的通讯消耗，那么此时在多个gpu上做不同head的并行计算，可能是比单个gpu同时做这些head的并行计算要快的。 但是如果考虑了通讯消耗，那么这个加速时间就有所折损 ，我们设1

建模 VS 实验结果

建模建完了，这下我们可以用真实的实验结果看看建模拟合效果怎么样了。在DistServe中，确实做了这样的一个实验，同时也提到了建模效果基本拟合的试验曲线的趋势（但是没有单独给出建模效果的趋势图） 。我们来看下实验效果：

如图，左侧为真实实验的结果。关于右侧的图，我这里有些迷惑，它似乎是在论证tp模型中取不同折损因子K值时对TTFT的影响（黄线），但是如果是这样，这个曲线应该是用建模结果表示的，但它现在却长得和真实实验结果一致。

但是anyway，回到左侧的真实实验结果上，我们不难发现，对于prefill阶段，在rate较小时，更适合用tp；在rate较大时，更适合用pp。也即prefill阶段在不同条件下，对并行方式有倾向性。

(2) decode阶段的并行策略

这里DistServe不再额外建模了， 但是它从另外一个角度做了模拟实验的结果 :

• latency：模型处理完一个请求的耗时
• throughput：单位时间内模型能处理的token数量

可以发现，对于decode阶段，随着gpu数量的增加，如果采用pp的方式，能产生更高的吞吐量（理由上面讲过，PP对batch的处理时流水线式的）；但是如果用tp，则能产生更低的延迟（即虽然单位时间内处理的tokens数量有限，但是处理单个请求的时间变快了）。

四、实操：代码实践中的prefill/decode并行策略

在上述的讲解中，我们通过实验的方式，比较直观感受到prefill和decode阶段对于不同并行策略的倾向性。我们还首次尝试通过数学建模的方式，来拟合在不同并行策略下系统的性能曲线 （不过正如前文所说，我们初次尝试的这个数学模型具有一些强假设，使得它不适合用在实践中。因此在实践中我们可能要考虑别的建模方式，我们马上来谈着点）。 在本节中，我们来具体看DistServe在代码实践中是如何为prefill和decode阶段分配最佳并行策略的。

在开始讲解之前，我们 先明确实操中DistServe的整体优化思路：

• （1） 对prefill和decode分别做goodput最大化，也就是对他们分别寻找最佳的并行策略（确定一个实例/一个模型副本要怎么分割）

• （2） 根据系统中的流量大小确定prefill/decode下一共要部署多少个实例，然后按照（1）中确定的并行方式做实例部署

• 在DistServe中，管这样的策略叫placement。

接下来， 我们再来看两种不同的硬件场景：

• 场景1：node间带宽十分高，这样我们可以忽略KV cache的传输时间 。我们使用algorithm1来寻找在这样的场景下对单个prefill/decode实例的最佳并行策略。值得说明的是，单node推理也属于该场景

• 场景2：node间的带宽并不高，不能忽略KV cache的传输时间 。这种情况下，我们尽量把prefill和decode实例对应的模型层分在同一个node中（也就是prefill和decode的pp维度必须相同）。这种是实践中更普遍的场景。

【本节根据论文伪代码及DistServe开源代码库撰写而成。https://github.com/LLMServe/DistServe/tree/main/simdistserve】

4.1 伪代码：寻找最佳并行策略