打造高性能大模型推理平台之Prefill、Decode分离系列（一）：微软新作SplitWise，通过将PD分离提高GPU的利用率

笑看风云

1. 背景

近期人们对大模型的使用量迅速增长，这为云商们带来了不小的挑战：要求其不断地部署更多更多更多的GPU。对于云商来说，更多的GPU意味着不菲的开销，钱包不开心；对于用户来说，由于算力跟不上造成自己的需求被拒绝，心情不美丽。
因此，如何能在已有的架构之下，更高效的利用资源，使得其在有限的硬件资源和电力条件下更快、更多、更好的处理请求将对云商和用户来说都意义非凡。
那么有一些学者们就发现，当前LLM推理过程中，限制其对资源更高效的利用的一个原因是：LLM推理过程中存在着两个截然不同的过程，那就是Prefill阶段和Decode阶段。在SplitWise这篇文章中分别把这两个阶段称为prompt phase 和 token-generation phase。



图1. LLM中的Prefill和Decode阶段

由图1可以看到，当一个Request进来时：

• 会首先经过Prompt phase，在这个阶段中，LLM处理所有用户的input或者prompt，计算出对应的KV Cache，其并行性可以充分利用GPU的算力，属于计算密集型。
  
• 接下来会进入Token generation phase，在这个阶段会顺序的产生一个个的token，每次访存只计算一个token，其对算力的要求并没，有那么大，主要受内存带宽的限制，属于内存密集型。
  

可以看到这两个阶段的特征完全不同，即便使用很好的batching技术，也无法解决两个如此明显不同阶段所带来的问题，比如：由于硬件资源利用不足，使得为用户提供服务将产生更高的花费。下图2可以展示了这两个不同阶段的一个例子。



图2. Prefill与Decode阶段的例子

2. 7个insight

因此针对这一问题，SplitWise主要针对7个insight做了详细分析，结论先来（详细分析在后~）：
1. 不同的推理任务中，Prefill和Decode的分布不同
2. 在P/D混合部署中，大部分时间消耗在Decode阶段，未充分利用到GPU的算力
3. 对大部分的请求，E2E时间消耗在Decode阶段
4. Prefill阶段应该限制Batch size从而避免影响性能，相反，Decode阶段应该增大Batch size来获得更高的吞吐
5. Prefill阶段算力是瓶颈，Decode阶段内存是瓶颈
6. Prefill阶段能充分使用算力，Decode阶段不能
7. Decode阶段可以使用性能较弱、成本较低的设备

2.1 Insight 1：Prompt和产生Token的数量




图3. Insight 1: Number of prompt and generated tokens

作者使用了LLama2-70B和BLOOM-176B两个模型来评估其参数量。并且针对两个不同类型的任务，分别是Coding和Conversation。对于Coding来说基本是代码补全任务，因此代表的是输入长、输出短的任务。上图的（a）可以看出，Coding和Conversation输入Prompt长度的中位数分别为1500个Tokens和1020个Tokens（具体数值论文中给出）。（b）则可以看出对于产生的token，Coding和Conversation分别有12个tokens和129个tokens。
对比（a）和（b）可以发现第一个结论：不同的推理任务中在Prefill和Decode的分布上非常的不同。

2.2 Insight 2：Batch的利用

首先我们来看一下几种不同的Batching



图4. 不同的Batching方法

（a）是以Request为级别的Batching方法，当有多个Request来了之后，可以按照时间顺序开始执行，一旦开始执行将完整执行Prefill和Decode阶段，只有全部执行完才能开始下一个Request。比如图（a）中Request 1和2同时到达，因此一起开始执行，后续就算来了3和4也要等到1、2完整执行完才能开始。
（b）在这个方法中，一个Batch里要么单纯的做Prefill，要么单纯做Decode。因为其认为Prefill会影响到TTFT，更重要，因此一个等待的Prefill可以抢占Decode，这种做法虽然会提高TTFT，但是也会造成TBT（ Time between tokens ，Decode阶段两个Token之间的时间）的尾延迟增高。
（c）中Prefill和Decode可以一起运行，这就会造成Prefill对Decode阶段产生干扰，本来可以很快执行完的Decode由于被Prefill干扰而拖慢。

了解清楚Batching的种类后，我们一起开看看作者对Batch利用的实验（使用mixed continuous batching）：



图5. Batch 利用

上图中的active tokens是说如果在Prefill阶段有100个tokens处于运行阶段，那么active tokens就是100。如果处于Decode阶段，由于一次只产生一个token，则active tokens就是1。
由上图可以看到，对于Conversation来说，有60%到70%的时间只运行了20个或者更少的tokens。对于Coding来说，有超过20%的时间仅仅运行了1个token。因此可以得出第二个结论：在mixed continuous batching中，大部分时间花费在Decode阶段。

2.3 Insight 3：时延




图6. Latency for BLOOM-176B and Llama2-70B on DGX-H100.

由上图可以看到Prefill和Decode时延都会随着token的数量增多而增多，而（c）更可以看出在E2E（end-to-end）的时延中，TTFT（Prefill阶段）时延仅占一小部分，尤其在Conversation中可以看到，只产生129个tokens所花的时间都比处理1020个Prefill的token时间多的多。论文中作者发现对于BLOOM-176B，Prefill阶段处理1500个token所花的时间与Decode阶段处理6个tokens所花的时间是一样的！
因此得出第三个结论：对于大多数请求，E2E的时延大部分花在Decode阶段。

2.4 Throughput




图7.  Impact of batching on the throughput for the two LLMs.

上图（a）可以看到，在Prompt phase（Prefill阶段），随着Batch的token数量增大，一开始的Throughput是会一起增大的，但是达到一定的值之后，throughtput不增反而开始下降。而与之相反，Token generation phase（Decode阶段）的throughtput则会随着Batch size的增大不断增大。
因此可以得出第四个结论：Prefill阶段应该限制Batch size从而避免影响性能，相反，Decode阶段应该增大Batch size来获得更高的吞吐。

2.5 内存利用




图8. Memory requirement with batching in prompt and token ph

由上图可以明显看到Prompt 和Token这两个阶段对内存需求非常不同。对Prompt phase 来说，内存需求不高，且相对稳定，而对于Token phase来说，内存需求会迅速增大。
因此得出第五个结论：Prefill阶段算力是瓶颈，Decode阶段内存是瓶颈

2.6 用电量




图9.  Maximum and mean power utilization varying the batching size.

由上图可以看到Prompt phase基本达到了TDP，算力可以拉满，而Token generation则相差甚远。



图10.  Impact of power cap on the prompt and token generation latency with the maximum batch size possible.

图10显示增加用电功率上限后的变化，可以看到如果将cap增加到一半（从700到350），对于Prompt phase时延增加会非常明显，而对于Token generation则基本没有变化。
综合上面的信息可以得到第六个结论：Prefill阶段能充分使用算力，Decode阶段不能。

2.7 不同的GPU硬件



上面的表中作者以H100为基准，验证在Coding和Conversation中用A100性能会下降多少，发现其实对TBT（也就是Decode阶段）影响较小，性能只是原来的0.7倍左右。
因此得出第七个结论：Decode阶段可以使用性能较弱、成本较低的设备。
以上就是作者详细分析的7个Insight。
3. SplitWise架构

有了以上的详细分析，验证了Prefill和Decode在各个方面的不同性质和影响，我们来看一下作者给出了什么样的解决方案。



图11.  High-level system diagram of Splitwise.

这张图展现了SplitWise的整体架构图，首先它有三个资源池，分别是Prompt Pool，Token Pool和Mixed Pool。在Prompt Pool中只做Prefill阶段，Token Pool中只做Decode阶段，实现了Prefill和Decode的分离部署。在Mixed Pool中以合并部署的方式执行Prefill和Decode混合的Batch，当请求负载变化剧烈、实例来不及伸缩或角色转换时，由Mixed Pool来执行混合的Batch。其次它还有两个级别的调度器，分别是cluster-level scheduler (CLS)和machine-level scheduler (MLS) 。
3.1 Cluster-level scheduling

• 在CLS中针对请求的路由策略采用“加入最短队列”的方法（Join the Shortest Queue (JSQ)）；在调度请求来的时候就决定后续要执行Prefill和Decode的实例。
  
• 在资源池中的机器管理方面，SplitWise在初始化将一个机器分配到那个资源池的时候根据请求的Token数和预估的负载来分配；资源池中的角色可以相互转换，但是由于要加载模型之类的操作，准换时间非常久，需要5分钟左右。
  
• 在Mixed Pool中，使用Mixed Batching，当在Prompt Pool或者Token Pool中queue 中的数量超过一个阈值的时候，就会在mixed pool中寻找机器，当mixed pool中的请求都执行完后，mixed pool中的机器将回到它原始的资源池中。
  

3.2 Machine-level scheduling

• 在MLS中，对于请求的策略采用FCFS，并且为了避免图7（a）中出现的Prompt阶段由于Batch 的token数太多导致的吞吐下降的问题，SplitWis将batching的prompt token总数限制在2048个。
  
• MLS会跟踪内存，当机器中的内存快要用完的时候，MLS会开始新排tokens。
  
• 在Mixed Pool中，给prompt的优先级更高（为了TTFT），因此可以发生抢占，但是为了避免饥饿问题，将会提高比较老的token的优先级，并且限制每个token可以抢占的次数。
  

3.3 KV-cache transfer

作者提出，他们给出的这种分离架构的方案引入的最大开销就是把KV Cache从Prefill阶段迁移到Decode阶段的开销。其为了减小这方面的开销，也做了一些优化：



图12. Optimizing KV-cache transfer in Splitwise.

（a）使用的是未经优化的线性KV Cache传输，那么所经历的时间就包括Prompt时间，等待KV Cache传输的时间和Token的时间。
（b）是SplitWise优化后的KV Cache传输，他们提出了Layer-wise的方法，以层为单位，在Prompt阶段，算完一层的KV Cache就开始传。
4. 实验效果

4.1 KV Cache的Latency

在刚才提到的KV Cache传输中，作者发现如果用Serialized的方法会给第2个token的时延增加64%，但是如果使用SplitWise优化后的per-layer的方法，仅给第2个token的时延增加16.5%.具体如下图所示：



图13. Overhead of the KV-cache transfer as the prompt size increases on A100s and H100s machines.

并且在这幅图中可以看到，当token数少的时候，反而是Serialized的方法更好，因为很少的token时就没必要花额外的开销分层，因此SplitWise设定当prompt sizes < 512 (in H100)时用Serialized，>512时用per-layer。
4.2 End-to-End 的Latency




图 14. Overhead of KV cache transfer on TTFT, E2E latency for coding trace for H100.

图14也印证了我们刚才的分析，每个token数对应3个柱子，左中右分别对应三种情况的E2E的Latency：采用per-layer方式、采用serialized方式，和用1台机器不做PD分离的对照组。可以看到在token数比较小的时候，确实serialized的方式比per-layer效果要好，随着token数逐渐增大，per-layer的效果变得比serialized更好。最终发现，采用Serialized的方式会给E2E的Latency增加3%，而用per-layer只增加0.8%。

4.3 ISO-Power 和 ISO-Cost



AA代表Prefill和Decode分离，分别使用A100；
HH代表Prefill和Decode分离，分别使用H100；
HA代表Prefill和Decode分离，Prefill使用H100，Decode使用A100；
这组实验为了证明insight 7 中Decode阶段可以使用相对便宜的设备。

ISO-Power：当固定用电量后，与 baseline-A100相比，Splitwise-AA在相同的用电量和花费的情况下吞吐提升了2.15倍之多。与baseline_H100相比，Splitwise-HA在相同的用电量，10%的更少花费下，吞吐提升1.18倍之多。
ISO-Cost：Splitwise-AA在相同的花费下表现效果最好，比Baseline-H100的吞吐提高1.4倍之多，使用了25%的更多电量，但是拥有了2倍的更多空

5. 参考文献

SplitWise  https://arxiv.org/abs/2311.18677

原文地址：https://zhuanlan.zhihu.com/p/701772045