OSDI24 DistServe：PD分离调度 论文阅读笔记

幸福侠侣

本笔记主要是对DistServe中大模型推理瓶颈和问题分析进行记录，有关具体分布式推理系统实现部分未做记录。

背景

大型语言模型（LLM）服务采用客户端-服务器架构：客户端提交一段文本作为请求，服务器在GPU上托管LLM，运行推理并将生成结果响应（或流式传输）回客户端。
现有系统的主要挑战在于它们必须在预填充阶段的第一个标记时间（TTFT）和解码阶段的每个输出标记时间（TPOT）之间做出权衡。
服务级别目标（SLO）。服务系统必须在满足这些SLO的同时，尽量降低与昂贵GPU相关的成本。
希望服务系统在每个GPU上最大化在满足SLO目标的情况下每秒处理的请求数量——即最大化每GPU的良好吞吐量。
1. 大模型推理：预填充（计算密集型）和解码（内存IO密集型）

当处理用户提示时，如果输入序列较长，预填充步骤往往受计算限制。
尽管解码阶段每步仅处理一个新标记，但它与预填充阶段一样，需要大量的I/O读取KV Cache，使其受到GPU内存带宽的限制。
在这两个阶段期间，在每个标记位置生成的中间状态（称为KV Cache）在后续解码步骤中再次需要。为了避免重新计算，它们被保存在GPU内存中。由于LLM权重和KV缓存的共享使用，大多数LLM推理引擎选择将预填充和解码阶段共同放置在GPU上，尽管它们的计算特性不同。
2. 推理优化：批处理与模型并行

批计算（效率与妥协）：当前都采用ORCA中的连续批处理，其特点在于PD两阶段任务都放在一起批处理。然而，这种方法在第一个标记时间（TTFT）和每个输出标记时间（TPOT）之间存在权衡。

一个改进的连续批处理变体（Sarathi-Serve）尝试通过将长预填充分割成块并将解码任务附加到这些块上来平衡TTFT和TPOT，但本质上，这只是以TTFT换取TPOT，无法消除两者间的干扰。同时长文本的Prefill分割后的成本会逐渐趋近于类似Decode阶段的访存密集型成本。

模型并行（提升性能的双刃剑）: 模型并行通常分为操作内并行和操作间并行.

• 操作内并行将计算密集型操作（如矩阵乘法）分割到多个GPU上，加速计算但引发大量通信。这减少了执行时间，从而特别降低了预填充阶段的TTFT，但需要高带宽的GPU间连接（如NVLINK）
  
• 操作间并行则将LLM层组织成阶段，每个阶段在一个GPU上运行形成流水线。尽管增加了由于阶段间通信导致的执行时间，但随着每个新增的GPU线性扩展系统的速率容量。
  

问题与挑战

预填充与解码干扰

图2显示，向解码请求批次中添加一个预填充任务会显著减慢两个过程的速度，导致TTFT和TPOT的明显增加。

• 具体来说，解码任务必须等待较长的预填充任务完成，从而延长TPOT；预填充时间越长，减速越严重（如图2(b)所示）。
  
• 在GPU负载已满的情况下，向预填充任务中添加解码任务也会增加预填充任务的完成时间（如图2中蓝色曲线所示）。
  

一种缓解方案是“块预填充搭载”的技术（chunked prefill with piggyback）：它提出将长预填充任务拆分为多个小块，并与少量解码任务批处理（即搭载技术）
但也有额外开销：

• 块大小设置得远低于能够饱和GPU的临界点，那么预填充任务的执行时间会更长，因为它与同一批次中的解码任务竞争GPU资源，无法单独利用GPU资源
  
• 如果增加块大小以接近饱和GPU，搭载的机会将减少，因为剩余的解码任务槽位有限
  
• 块预填充会显著增加预填充任务的内存访问量：如果将一个预填充任务分为N个相等的块，则需要总共加载N + (N −1) +...+1 = O(N²)个KV缓存块，而在非块情况下仅需O(N)。随着上下文长度的增加，这一开销将会增加。
  

无效的调度

将预填充和解码任务拆分开并顺序调度并不能缓解干扰。解码任务可能因等待正在进行的预填充任务而经历更长的排队延迟。此外，专门用于解码的批次往往导致GPU利用率不足。优先处理某一阶段的任务会对另一阶段的延迟产生不利影响，使优先级调度无效。
资源和并行性耦合

将预填充和解码阶段共置在同一GPU上，不可避免地共享其资源和并行性设置。然而，每个阶段都有其独特的计算特性和延迟要求，需要更为异构的资源分配。

• 预填充阶段往往受计算能力限制，受益于更多的内部操作并行性，以减少执行时间以满足TTFT的严格SLO。
  
• 相比之下，解码阶段的最佳并行性配置取决于运行的批次大小。
  

在现有系统中，由于耦合的存在，资源分配和并行性计划通常为了满足TTFT和TPOT中更苛刻的要求，而这可能对另一个阶段并不理想。这往往导致资源的过度配置，以满足两个SLO。
动机

分离预填充和解码阶段自然解决了两者之间的干扰，使每个阶段能够专注于其优化目标——TTFT或TPOT。每种实例类型可以采用不同的资源和并行性策略，以满足各种延迟要求。通过调整分配给这两种实例的GPU数量和并行性，我们可以最大限度地提高整个系统的单设备良品率，避免过度配置，最终实现减少每个查询的成本，同时确保服务质量。
批处理策略

预填充阶段通常是计算密集型的。图3(a)中可以看到prefill阶段很容易达到Compute Bound。如130亿参数的LLM来说，处理512个tokens的单个序列就能完全占用一块A100 GPU

• 对于预填充实例，有必要预先对特定的LLM和GPU进行配置文件分析，以确定一个关键的输入长度阈值，即Lm，超过该阈值后，预填充阶段就会变得计算密集。只有当计划请求的输入长度低于Lm时，才应考虑进行批处理。在实际操作中，用户提示的平均长度通常在几百个tokens左右，预填充实例的批处理大小一般保持较小。
  

解码阶段通常是访存密集型。单个解码任务严重依赖带宽，因此批处理是避免GPU利用率低（从而提高每个GPU的吞吐量）的关键



并行策略

这里使用建模方式来探究张量并行还是流水线并行更好
Prefill阶段：

• 图4(a)可见：更严格的SLO使得操作内并行更具优势，因为它能减少执行时间。
  
• K的值取决于输入长度、模型架构、通信带宽和布局等因素。如图4(b)所示，K值的降低显著减少了操作内并行的效率。（后续章节该工作开发算法来考虑这些参数优化资源和并行性配置）
  

Decoding阶段：

• 随着解码批处理大小的不断增加，解码计算开始类似于预填充阶段。
  
• 如图5所示：操作内并行（intra-op parallelism）虽然可以减少延迟，但回报递减，因为通信开销和分区后利用率降低。操作外并行（inter-op parallelism）几乎可以线性地扩大吞吐量。
  

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