速读论文-ISCA'24 SplitWise 解耦算力依赖性独立优化
叠甲️：这个系列以一些更短平快的风格论文读后感博客，囫囵吞枣不求甚解。架构文章往往涉及建模仿真或实测数据，完整复现工程复杂度高，而不复现又学不懂。这个系列找到简化角度敏捷建模学习，尽可能以几个问题、几句总结切入概括文章。因为简化模型添加许多假设，数据结果大概率和原论文存在差距，更多侧重原文 motivation 思路和此类问题思考方式。
都 6202 年了，还在看 PD 分离呢。虽然 PD 常常口头谈及，仔细想想细节问题还是不能一下子答上来，故温故文章。提及 PD 分离第一印象往往是 roofline 图， perfill 高并发偏 compute bound，decode 低并发偏 memory bound。但另一个印象是，即使不做 compute-bandwidth 异构硬件，在集群层面划分 PD 也能有收益。相比将 PD 放在一台机器上，PD 划分不仅添加了额外 KV Cache 传输开销，也没有改变硬件计算-带宽比例，为什么能够有提升呢？
一言以避之，PD 分离解开 PD 优化的耦合关系，进而解放两个阶段相互制约的硬件优化条件，让每个阶段能够各自收敛到负载曲线更优的解。 而异构硬件则是在这个前提上，将负载曲线差异性拉得更加夸张，提高解的上限。用两个关键词概括 PD 分离：（1）解耦合；（2）负载曲线。下文将以同构硬件 PD 分离做讨论,以 Splitwise 论文数据切入展开[1]。本篇 blog 相关代码已开源[2]。
建模假设
为了建模方便，做出以下建模假设：
- 机器粒度： 这里将能够独立推理完整模型的硬件单元叫做机器，一台机器可能有多个计算卡组成，机器内部以某种并行策略部署。本文后续最小粒度仅控制在机器层面讨论。
- 理想数据传输： 不考虑 PD 之间传输 KV Cache 数据开销；
- 单轮 request： 不考虑多轮 request KV Cache 积累状态量，仅以单轮 request 作为切入；
- decode 长度无关容量曲线： decode 阶段 batch 数量和容量的曲线不考虑 decode length 的变化，这个假设十分不合理，但为了避免引入负载工况从原始论文提取数据只好如此，或许可以理解为这种假设下的结论代表某种期望平均值。
- decode 阶段 beam=1。
Baseline PD Co-located Policy: 流水线的耦合优化关系
Prefill -> Decode 是流水线阶段关系，一次请求（request）以 prefill 计算开始，以 decode 输出结尾。流水线意味着短版效应，系统开销取决于最差的环节。不过短板效应有两种表现，如果在同一台机器上划分资源，因为兼顾另一者两者都不能取到资源-成本较优；如果按空间划分资源，则是较快者会出现闲置等待较慢者完成。
Baseline 模式是 PD 在同一台机器上完成，那么 PD 阶段之间存在什么依赖？又耦合竞争什么资源呢？
优化关系：batch- 容量、吞吐曲线
核心是 Prefill 和 Decode Phase 对 Batching 的 data scaling 趋势不一致。可以抽象为容量(Capacity)和吞吐(Throughput) 的变化趋势 \(C_{p}(B_p);C_{d}(B_d);T_p(B_p);T_d(B_d)\)。其中对于 prefill \(B_p=nL_p\) ；对于 decode \(B_d=n\) （beam=1）。其中 prefill 和 decode 长度分别是 \(L_p, L_d\)，同时请求数并发数量为 \(n\)。
如图是容量和吞吐的曲线关系。因为 decode 容量需求不仅仅和 batch 相关，还和 prefill 长度相关，所以拟合的时候是一群散点。以及本身 batched token 含义不同，PD 曲线在同 x 坐标对比没有意义。这里大致可见，prefill 和 decode 随着 batched token 上升，容量需求都会增加；但吞吐上 decode 随着 batch 数量变化显著，而 prefill 几乎不变，token 已经打满饱和。
资源竞争：P D 在机器层面竞争存储容量
其中 \(C\) 是机器容量， \(C_w\) 是模型权重容量。 P 和 D 处理并发请求数量相等 \(n_p=n_d=n\)
流水线依赖：通过分配时间比例保证 RPS 平衡
在同一台机器划分 PD，对于 PD 需要在 机器层面取得平衡。所以其次平均请求处理速率 （Request per Second, RPS）相等，一台机器划分不同时间段处理 prefill 和 decode。假设 \(\alpha\in (0,1)\) 的时间比例分配给 prefill， 一共有 \(X\) 台机器 。则理想各个阶段吞吐率为 \((XT_p(nL_p),XT_d(n))\)：
prefill 和 decode 的划分 \(\alpha\) 应该满足：
其实这里应该也存在分 small batch 多次处理的空间，比如 prefill 分 p 次处理，decode 分 d 次处理，RPS 平衡公式则变成：
但从吞吐曲线更像是一个凸曲线，这样划分不如单次处理完，batch 越大越好了。
优化目标：两个阶段的 token 吞吐率
对于集群， prefill 和 decode 平均吞吐速率为 \((\alpha XT_p(nL_p),(1-\alpha) XT_d(n))\)。PD 的吞吐曲线位置和 batch 数量有关，而在同一台机器上，batch 数量受到机器粒度请求数相等以及容量竞争的限制，导致 P 和 D 互相影响了对方的优化空间。
PD Split Policy： 静态划分众口难调
接下来考虑 PD 分离，假设 \(x\) 台机器只用于 decode，\(X-x\) 台机器用于 prefill。
- 优化关系： 同构硬件，仍然是前文曲线不变；
- 容量资源竞争： 二者不在机器层面产生竞争，理想各自用满机器容量，\(n_d=C_d^{-1}(C)\)，\(n_pL_p=C_p^{-1}(C)\)。集群层面处理并发数量相等，\(xn_d=(X-x)n_p\)。
- 流水线依赖： 通过分配机器在集群层面 RPS 平衡。\((X-x)\frac{T_p(n_pL_p)}{L_p} = x\frac{T_d(n_d)}{L_d}\)
- 优化目标： 平均吞吐率为 \(((X-x)T_p(n_pL_p),xT_d(n_d))\)。
看里出现了两个求解分配比例 \(x\) 的方程，导致 \(x\) 无解。这是因为在 Co-located Policy ，存在两个可变变量 \(n\) 和 \(\alpha\) ，进而满足容量约束和 RPS 平衡约束。这里做了各自用满机器容量的假设，导致消去了 \(n\) 的自由度，所以更加准确的是，要么用满容量，但允许 PD 集群有 under-utilization 限制状态；要么用不满容量，但 PD 满 utilization，空间静态划分策略导致难以同时满足两个平衡条件，但这也提供了 mixed pool 的优化空间，将一部分限制的 decode 资源用以执行 prefill （decode 可以插入 prefill ，而 prefill 不可以插入 decode 主要是根据 PD 的依赖顺序和多轮交互 KV cache 传输角度考虑的），不过 mixed pool 就不在本文讨论了。
故 PD Split Policy 应该分为两种子策略：
- Capacity Utilization-first： \(n_d=C_d^{-1}(C)\)，\(n_pL_p=C_p^{-1}(C)\), \(x\) 满足 \(xn_d=(X-x)n_p\) ，存在利用率因子 \(\alpha_p\) 及 \(\alpha_d\) （其中一个取到1），利用率通过 \(\alpha_p(X-x)\frac{T_p(n_pL_p)}{L_p} = \alpha_d x\frac{T_d(n_d)}{L_d}\) 求解，优化目标平均吞吐为 \((\alpha_p(X-x)T_p(n_pL_p),\alpha_d xT_d(n_d))\)；
- Time Utilization-first: \(n_d\) 和 \(n_p\) 一个用尽容量，\(n_d=C_d^{-1}(C)\) 或 \(n_pL_p=C_p^{-1}(C)\)；另一个则满足容量约束，\(C_d(n_d)<C\) 或 \(C_p(n_pL_p)<C\) ，满足 \(xn_d=(X-x)n_p\) 以及 \((X-x)\frac{T_p(n_pL_p)}{L_p} = x\frac{T_d(n_d)}{L_d}\)，优化目标平均吞吐为 \(((X-x)T_p(n_pL_p), xT_d(n_d))\)。
粗糙实验数据分析
带入原文提取的拟合公式以及硬件数据，分析不同 workload 负载（更改 prefill 和 decode 长度比例，具体 Token 长度数值从 Kairous 论文提取[3]）；研究不同集群规格数量下，吞吐率、最优 PD 分离策略以及分离比例关系。
实验结果表明，在同构硬件集群下
- 吞吐率： 大部分 workload （LPLD、LPHD、HPLD）PD 分离体现微弱提升或者负面提升；而 HPHD 有大致 ~1.5x 左右提升；
- 集群数量： 集群数量节点越多，分离优势越大（毕竟调度空间大了，合理）；
- 策略选用： 大部分策略优先满足 time utilization 并打满 decode batch （毕竟 batch 对 decode 提升收益更大）；
- 分配比例： 大部分 workload (LPLD、HPLD、HPHD) 将更多节点分配给 prefill，LPHD将更多节点分配给 decode。
思考：系统重构能力和成本的 trade-off
首先回答博客开始的问题，同构硬件 PD 分离虽然没有改变硬件计算-带宽的能力，但解耦开了 PD 分离的容量约束进而提升系统效率。计算-带宽-容量三者才是设计负载最基础抽象。通过解耦开资源依赖/竞争，并赋予硬件重构能力，进而让每个阶段独立优化，各自达到局部更优进而达到全局更优。
以往学习更多涉及 accelerator ，因此应用局限在 chiplet 或是架构内部等重构方式。往往因为容量不够/片外带宽不够感到无计可施或者大炮打蚊子。不同的硬件配置方法，存储-带宽-计算不同的参数、设计异构硬件 SKE 配合、不同封装技术、更改集群比例数量他们的成本和上限收益能有多少呢？如何推导他们的 cost model 快速做决策？架构层面解开一个锁可以有很多把钥匙，更关键的或许是分析能力和成本的 trade-off，根据问题选择最合适的钥匙，让问题轻盈丝滑地落地。或许也正是在不同层次简化问题，快速找到较优解，才能带来 AI 基础设施的快速演进。
浙公网安备 33010602011771号