PD分离思考

营养师

前言

PD 分离出来也有一阵子时间了，之前一直觉得主要问题在 KV Cache 传输上，并且业内似乎落地成功的不多于是没有过多思考，碰巧前阵子 mooncake 正在开源他们的分离架构，于是这篇文章对 PD 分离进行一个进一步的思考。
融合推理存在的问题

PD 分离主要是考虑到了 LLM 的 prefill 和 decode 的两个阶段的特性不同，prefill 阶段是 compute bound，decode 阶段是 memory bound，prefill 阶段的能力我们用 TTFT 首 token 时延来衡量，decode 的能力我们用 TPOT 生成每个 token 的时间来衡量。
但是在同一张卡上做 prefill 和 decode 会出现问题，在机器的算力等条件固定的情况下，你增加 bsz，prefill 阶段机器到算力瓶颈了，反而影响 TTFT，你减小 bsz，decode 阶段又是访存瓶颈的，decode 阶段可以比 prefill 阶段承载更大的 bsz。
那么问题来了，到底要不要增大 bsz？
vLLM 中的调度策略是：
优先调度 prefill，然后调度 decode，但是会设置一个阈值防止 decode 一直等待。
也就是优先确保 TTFT，然后设置一个阈值来保证 TPOT 不会太差。很明显，这只是一个折中的方法。
PD 分离的优势

有了 PD 分离之后，我们可以把 prefill 阶段放在 H800 这样的算力高的机器，decode 阶段放在 H20 这样算力低的机器但是访存能力不会差太多的机器（毕竟显卡更新换代过程中算力增长是遥遥领先访存能力增长的），这样我们的如何 bsz 如何均衡的问题似乎可以得到解决，不同机器只负责一个阶段，bsz 也只需要根据你这个阶段的特性来设置就好了。并且这样还能节省机器费用利用上了 H20 这样的便宜的卡，何乐而不为呢？
PD 分离的缺点

但是 PD 分离有个很重要的问题，增加了通信和网络传输的成本，如果是卡间分离那么会增加通信的成本，如果是不同机器上进行分离那么就会增加网络传输 KV Cache 的成本，毕竟 KV Cache 很大(十几甚至几十个G呢)。
<hr/>如果错误，请不吝指正～

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