SageMaker HyperPod 上实现 LLM 推理的预填充与解码分离架构
大语言模型推理中的预填充（Prefill）与解码（Decode）共享同一 GPU 时，长提示会阻塞所有并发请求的令牌生成。分离式预填充与解码（DPD） 通过将两个阶段分别运行在独立的 GPU 池中，并借助弹性结构适配器（EFA）与远程直接内存访问（RDMA）连接，消除了这种干扰。
LLM 推理包含两个本质不同的阶段：预填充是计算密集型，它并行处理整个输入提示以生成初始键值（KV）缓存；解码则是内存密集型，每次生成一个令牌，需要大量内存带宽来访问模型权重和不断增长的 KV 缓存。通过将两者分离为专用引擎，可以为每个阶段分配不同的并行策略。这种分离允许独立调整首令牌时间（TTFT） 和令牌间延迟（ITL），比块状预填充调优更可靠地控制尾部延迟，并防止长上下文预填充阻塞正在进行的解码请求。
vLLM 通过连续批处理和 PagedAttention 提高了单节点效率。然而，大规模部署的组织在编排多节点部署和优化路由时仍面临挑战。本文展示了如何使用 SageMaker HyperPod Inference Operator 在 Amazon SageMaker HyperPod 上实现基于 vLLM 的 DPD。
何时使用分离式推理
分离预填充与解码在以下场景中效果最为显著：长上下文、高并发流式工作负载，例如聊天助手、Agent 管线、文档分析端点和带有大量检索上下文的检索增强生成（RAG）。在这些情况下，单个长提示在共享 GPU 上会导致其他所有请求的解码停滞，造成每令牌延迟尖峰，而 DPD 通过架构设计消除了这一问题。
考虑使用 DPD 的场景：
- 输入提示经常超过 4096 个令牌
- 存在多个并发用户或请求
- 需要流式响应且令牌交付一致性至关重要
- 混合流量中同时包含长提示和短提示
当 GPU 争用并非实际问题时，共存部署是更简单的选择：批处理或离线工作负载（优化 TTFT）、低并发部署、或仅短提示流量。在路由阈值以下，通过 EFA RDMA 传输 KV 缓存的固定成本超过了隔离解码带来的收益。DPD 路由器会将此类请求直接发送给解码器。
架构与实现
DPD 架构由三个核心组件构成：预填充池、解码池和智能路由器。预填充池负责处理提示并生成 KV 缓存，解码池利用接收到的 KV 缓存逐令牌生成输出，路由器则根据提示长度、当前负载和延迟目标决定请求的路由方式。
在 SageMaker HyperPod 上实现时，利用 HyperPod Inference Operator 可以简化多节点部署的编排。EFA RDMA 提供了低延迟、高带宽的节点间通信，使得 KV 缓存在预填充和解码池之间高效传输成为可能。
性能优化与权衡
分离架构带来的主要收益包括：
- 消除干扰：长提示不再阻塞其他请求的解码，尾部延迟显著降低
- 独立优化：可为预填充和解码分别调整并行策略，例如预填充使用更大的张量并行度，解码使用流水线并行
- 资源效率：根据工作负载动态调整预填充和解码池的大小，避免资源浪费
然而，也存在一些权衡：
- KV 缓存传输开销：通过 EFA 传输 KV 缓存会引入固定延迟，对于短提示可能得不偿失
- 复杂性增加：需要管理两个独立的 GPU 池和智能路由逻辑
- 路由决策延迟：路由器本身可能成为瓶颈，需要高效的路由算法
总结
DPD 是应对现代 LLM 推理中长上下文和高并发挑战的有效架构。通过在 SageMaker HyperPod 上结合 vLLM 和 EFA RDMA，组织可以构建可扩展、低延迟的推理服务。对于以流式响应为核心的应用，DPD 提供了一条清晰的路径，在保持高吞吐的同时实现稳定的每令牌延迟。
