【自用】AI infra相关：PD分离

在LLM推理过程中prefill阶段和decode阶段具有截然不同的计算特性prefill阶段需要并行处理整个输入序列来生成首个token属于计算密集型操作。decode阶段逐个生成后续token需要频繁访问kv cache属于内存密集型操作。传统的continuous batching将两个阶段混合处理导致互相干扰难以同时满足TTFT首token延迟和TPOTtoken间延迟的严格要求。为了解决这一问题PD分离框架应运而生通过将prefill和decode分配到不同的gpu实例上针对各自特性进行专门优化。这种分离式设计不仅消除了阶段间的干扰还能显著提升系统的有效吞吐量Goodput为大规模LLM服务提供了更优的解决方案。一、吞吐量Throughput) vs 有效吞吐量Goodput)目前大多数LLM都以吞吐量作为主要性能指标——即单位时间内处理的请求数RPS或生成的token数。实际上下游应用的类型多种多样它们在用户体验上的延迟需求各异因此需要满足的服务等级目标SLO也存在显著差异。大模型服务中最常用的SLO包括TTFTtime to first token首token响应延迟直接影响用户的等待体验。TPOTtime per output token衡量两个连续生成的token之间的平均延迟决定交互的流畅程度。例如实时聊天机器人更关注低TTFT以保证响应及时而TPOT只需快于人类阅读速度约250词/分钟即可相反分档摘要则更强调TPOT一遍更快地产生完整摘要。单纯依赖吞吐量作为指标并不能反映延迟表现系统看似处理了大量请求但其中不少未能满足SLO最终呈现给用户的仍是不理想的服务体验。Throughput通常指系统单位时间内处理的token数或请求数。很多工作把『提高吞吐量』作为主要优化目标但在实际场景下这并不直接代表用户体验。Goodput指系统在满足延迟约束如TTFT/TPOT SLO的前提下每秒完成的有效请求数。与单纯的吞吐量相比Goodput是更优的衡量指标因为它能够体现请求在满足SLO情况下的吞吐水平从而同时反映成本效益和服务质量。Goodput(P90 TTFT 200ms 且P90 TPOT 50ms)表示在至少90%的请求同时满足TTFT 200ms和TPOT 50ms的条件下系统所能维持的最大每秒请求数。eg:某应用的吞吐量为10 RPS每秒请求数但由于延迟约束的限制只有3 RPS的请求满足SLO因此该系统的Goodput仅为3 RPS。二、prefill和decode共置导致干扰在LLM服务中请求的生命周期通常包含两个阶段prefill生成首个token和decode逐步生成后续token。大多数现有系统如vllm,TensorRT-LLM)采用continuous batching集数将prefill和decode混合在一起统一批处理。这种方式确实能够提升整体吞吐量但由于两者计算特性和SLO目标差异显著将它们共置在同一gpu上往往不理想。如下图所示continuous batching会带来明显干扰。当prefill和decode被放在同一批次时decode请求的延迟TPOT)会被显著拉长而prefill请求的首token延迟TTFT也会有所增加。图中展示了三种不同的执行方式1PnD棕色1个prefill与n个decode混合批处理。nD蓝色仅包含decode请求的批处理。prefill-only红色虚线仅运行prefill请求的延迟。在prompt长度为128时相比仅包含decode的请求延迟增加约1.8倍而当prompt长度为1024时干扰效应显著放大decode延迟提升至12.6倍。由于这种干扰当服务必须同时满足TTFT和TPOT的SLO时系统往往需要进行资源的过度配置才能达到延迟目标尤其是在任一SLO要求较严格的情况下。三、PD分离的整体思路直观的思路将prefill和decode分离到不同的gpu上并为每个阶段定制并行策略。这自然解决了前面提到的两个问题没有干扰prefill和decode各自独立运行更快地完成计算也更容易满足各自的SLO。资源分配和并行策略解耦可以针对prefill和decode分别优化。当一个请求到达系统时它会先被分配到prefill worker完成prefill阶段随后系统将其中间状态主要是kv cache迁移到decode worker并执行多步decode以生成后续token当生成完成后请求才会离开系统。四、分离式推理架构的优化方向4.1算力与存储prefill阶段拥有计算受限的性质compute-bound)特别是在请求流量较大用户的prompt也较长的情况下。prefill阶段算完kv cache并发给decode阶段后理论上prefill就不再需要这个kv cache了当然也可以采用LRU等策略对kv cache的保存做管理而不是一股脑地清除decode阶段拥有内存受限的性质memory-bound)因为逐个token的生成方式decode要频繁从内存中读取kv cache同时也意味着它需要尽可能保存kv cache。因此在分离式框架下计算和存储可以朝着两个独立的方向做优化。4.2batching策略prefill阶段随着batch_size的增加吞吐量的提升很快趋于平缓这是因为prefill属于compute-bound当batch中的总token数超过一定规模后GPU的计算能力已经被完全吃满再增加请求只会延长整体处理时间而不会带来明显的吞吐提升。decode阶段随着batch_size的增加吞吐量的增长趋势越来越显著。这是因为decode阶段是memory-bound即相比于计算读写数据的时间要更多所以在decode阶段中如果我们能提升batch_size就能把计算强度提起来吞吐量就上去了。在分离架构下我们可以针对prefill和decode的特性对Batching策略分别进行优化具体来说对于prefill实例需要事先结合特定的llm和gpu做性能分析找出输入长度的临界点——一旦超过这个点prefill就会进入compute-bound此时增加batch_size只会拖慢整体处理速度。在实际应用中用户的prompt往往已有数百个tokens因此prefill的batch_size通常保持较小。相对的decode阶段更适合采用较大的batch_size已充分提升gpu利用率和整体吞吐。4.3并行策略由于prefill和decode具有不同的计算模式和延迟目标这两个阶段的最佳并行策略不相同。例如当TTFT要求严格而TPOT要求相对宽松时prefill更适合采用张量并行来满足低延迟而decode则通常采用数据并行或流水线并行来提升吞吐。五、kv cache传输PD分离带来的代价是需要在prefill和decode的gpu之间传输中间状态即kv cache。5.1kv cache传输开销初看之下kv cache是llm推理中巨大的内存开销而gpu之间kv cache的传输似乎会成为瓶颈。然而DistServe的论文中展示了相反的结果通过合理的放置kv cache的传输开销可以被有效地最小化甚至低于一次decode步骤的时间这得益于当今高速互联网络。假设在gpu之间使用8通道PCle 5.0x16每条链路64GB/s作为节点内互联。对于一个包含2048tokens的请求在服务OPT-175B时传输kv cache的延迟可以估算如下Latency2048 tokens * (4.5MB/token)/(64GB/s * 8)17.6 ms对于 OPT-175B延迟小于单次 decode 步骤在 A100 上约为 30-50 毫秒。对于更大的模型、更长的序列或更先进的网络例如带宽为 600GB/s 的 A100-NVLink如下图所示与单次 decoe 步骤相比KV cache 传输相关的相对开销变得不那么显著。总之通过精心安排 prefill 和 decode 工作节点以利用高带宽网络可以有效隐藏 KV cache 传输的开销。5.2kv cache传输方式目前kv cache的传输主要有两种方式中心存储和点对点P2P当然在实际系统中也可能采用两种结合的混合方案。中心存储建立一个跨设备的kv cache由它统一管理kv cache的增删查和传递等操作。prefill和decode实例只需与这个kv store交互负责写入或读取数据。p2p每个实例独立管理自己的存储。例如一个prefill实例完成计算后会直接与目标decode实例建立通信将kv cache传过去不依赖统一的中介。两种方式各有优劣中心存储更适合构建大规模集群能充分利用多种存储介质和传输通道并提升计算结果的复用效率但在某些场景下性能可能受限同时系统维护成本高。p2p架构很简单性能表现通常更好但在扩展性和链路稳定性方面会面临挑战。5.3kv cache传输的网络堆栈现有的物理数据链路可以分为3类Direct即GPU之间通过高速告诉直连链路如NVLink或HCCS相互连接。在这种情况下可以利用底层的内存拷贝原语或集体通信库来完成数据传输。Direct-NIC即GPU通过其配套的网卡NIC进行通信。在这里可以使用定制化的库通过PCle和以太网或InfiniBand进行数据传输。Indirect即当GPU之间没有直接链路时必须通过其CPU的DRAM中转数据从而带来额外的内存拷贝开销。5.4Kv cache传输粒度请求级等到prefill阶段完成后将kv cache一次性传输。这种方式的好处是能够减少网络传播次数因为每次传输的数据量更大从而降低了通信开销。然而当kv cache大小较大时会影响TTFT。层级spilitwise通过在prefill阶段的计算与kv cache传输之间实现重叠来优化性能。每一层计算完成后都会异步传输该层的kv cache同时继续执行下一层的计算从而降低传输开销。层级传输还能带来额外优势例如更早启动decode阶段以及更早释放prefill端的内存。层级kv cache传输与下一层的prefill计算并行进行这需要逐层的细粒度同步以确保正确性因此可能带来性能干扰并增加TTFT尤其是在小prompt场景下不过对于小prompt来说kv cache的总体规模很小不需要层级传输来隐藏延迟。由于在计算开始时批次中的token数是已知的splitwise会选择最合适的kv cache传输方式小prompt使用序列化传输而大prompt使用层级传输。块级TetrilInfer在PD分离的基础上还会将输入的prompt划分为固定大小的chunk以便让GPU始终运行在接近计算饱和的状态。5.5vllm的PD分离vllm提供了kv connector作为管理实例间kv cache交换的抽象层它提供统一接口来实现kv cache的保存、加载与传输使不同的vllm实例如prefill和decode实例能够高效共享计算结果。通过实现这一接口各类connector例如通过文件系统的SharedStorageConnector、通过网络的NixlConnector等提供了灵活的kv cache传输方案从而支持PD分离等高级功能。KVConnectorBase_V1是所有connector的基类他是一个抽象类定义了一下APIscheduler侧方法build_connector_meta构建元数据shceduler告诉worker需要保持/加载哪些kv cache。get_num_new_matched_tokens获取远端已计算的kv cache的token数量update_state_after_allocblock开辟后更新connector的状态。worker侧方法start_load_kv从connector buffer加载kv cache消费端调用。wait_for_layer_load阻塞直到指定层加载结束消费端调用。save_kv_layer将vllm的kv buffer中某一层的kv cache保存到connector buffer中生产端调用。wait_for_save阻塞直到所有保存操作完成生产端调用。vllm v1中connector有两个执行角色scheduler_connector和worker_connector分别在scheduler线程和worker线程中执行。shceduler负责指挥worker进行kv cache的传递两者之间的信息桥梁是元数据KVConnectorMetadataworker通过metadata知道哪些kv值需要从远端加载。当前vllm支持5种类型的connector分别为SharedStorageConnectorvllm中最简单的kv connector实现通过共享文件系统如本地磁盘或NFS在prefill和decode之间传递kv cache使用MD5哈希生成唯一文件名来存储和检索每个请求的kv cacheprefill实例将每层的kv cache序列化为SafeTensors格式保存到指定路径decode实例根据相同的token_ids计算哈希值找到对应文件并加载整个过程没有显式的网络传输完全依赖文件系统的读写操作。P2pNcclConnector基于NCCLNVIDIA Collective Communications Library实现的高性能kv connector通过NCCL的send/recv原语实现kv cache在不同gpu之间的点对点传输避免了文件系统的开销。NixlConnector使用NIXLNVIDIA Inference Xfer Library库来加速GPU之间以及异构内存与存储之间的kv cache传输。LMCacheConnectorV1通过与LMCache集成实现kv cache的外部存储与检索支持多种存储后端如CPU内存、本地文件系统、Redis、InfiniStore等。LMCache通过重用缓存的kv cache来减少推理时间消除冗余计算适用于跨请求或跨会话的kv cache共享场景。MultiConnector允许同时使用多个kv connector来实现kv cache的传输它的核心逻辑是从第一个能提供可用token的connector加载kv cache但会向所有connector保存数据。multiconnector适用于需要同时向多个存储后端保存Kv cache的场景比如同时保存到本地存储或远程存储提供数据冗余和可靠性保障。