文章

分离的预填和解码

高速几何数据粒子迅速飞驰而过

为了从输入提示生成输出令牌,LLM 推理被分成两个阶段:预填充和解码。预填充在输入令牌上运行,填充 KV 缓存,然后进入逐个生成令牌的解码阶段。

尽管单个解码步骤通常运行数十毫秒,但预填充所需的时间明显更长。如果在相同设备上运行,将预填充与解码混合会降低解码性能。在本文中,我们探讨了一种通过分散预填充和解码,并在不同设备上运行的成熟解决方案,以最大化预填充吞吐量和解码延迟。

预填充与解码性能

在典型的 LLM 服务引擎中,批处理调度器会选择在模型的每个执行步骤中处理的请求。当在单个设备或节点上运行时,预填充和解码请求被一起批处理。伴随序列长度聚合的注意力成本,随着 KV 缓存中的条目长度(kv_len)的增加而增加。解码请求通常只转发单个令牌(qo_len=1),通过其他独立于序列的令牌的层以最低成本运行。预填充请求以极高的成本通过高密度层(大的 qo_len)转发数千或数万个令牌。

前向传输的延迟更多地受到通过密集层传递的独立令牌数量(qo_len)的影响,而不是通过注意力从 KV 缓存中检索的令牌数量(kv_len)。注意力可以跨请求数量和与序列长度成比例的 kv_len 并行化,从而实现良好的资源利用。预填充是计算绑定的:由于 qo_len 高,GEMM 内核可以沿 M 维分配足够的块以充分利用现代 GPU 的计算能力。解码是内存绑定的:由于通常批量较小,沿 M 输入数量通常很小,仅足够一个块。虽然 Split-K GEMM 内核可以提高低令牌批处理大小的 SM 利用率,但缓存和矩阵乘法单元通常未被充分利用。

当混合在一起时,包含预填充请求的批处理在前向传输过程中会导致更高的延迟,从而对整个实例的解码吞吐量产生负面影响。尽管将预填充请求与解码请求混合或采用分块预填充可以略微提高解码性能,但很难保持足够的预填充吞吐量以在一个实例上处理足够的请求以最大化解码吞吐量。在大型模型的情况下,为了保持解码的大批量,预填充必须足够频繁地执行,这显著增加了平均延迟并导致输出卡顿。

通过使用一组单独的节点来执行预填充和解码可以解决这些问题。通过将一个预填充节点与多个解码节点关联,可以调度足够的请求以最大化吞吐量,并在解码节点上保持足够多的并发请求以同样最大化解码吞吐量。预填充节点填充 KV 缓存,然后将其传输到解码节点。由于解码器无需暂停进行预填充,延迟变得更具确定性,因为活跃请求增加的 kv_len 的整体影响变得不那么明显。代价是 TTFT(首次令牌时间)的增加,因为在网络上传输 KV 缓存可能需要数十到数百毫秒。

KV Messenger

在 Perplexity,我们的分散预填充和解码实现是基于与 LLM 引擎交互的 KV 信使,协调从预填充节点到解码节点的 KV 缓存传输。在预填充端,信使接受来自解码节点的请求,将其交给批处理调度器并跟踪前向传输执行,以尽可能低的延迟派发 KV 缓存。在解码端,分配不可逐出页面后,信使阻止请求被调度进行解码,直到接收到 KV 缓存和解码器上下文传输完成的通知。

分散预填充需要高吞吐量、低延迟的连接,因此我们的实现专为 RDMA 量身定制,支持 EFA 和 ConnectX 网络接口控制器(NIC)。KV Messenger 构建于 libfabric 之上,使用我们的 fabric-lib 包装器提供远程直接内存访问(RDMA)原语的高层低延迟抽象,实现高效的页面和元数据传输以及低延迟信号。在后台,fabric-lib 协调 GPU 及其直接连接的 NIC,将数据从预填充节点复制到解码节点。

在接收到预填充节点分配了一组相应的源 KV 页面后,使用其本地引擎调度预填充请求。为了最小化延迟,传输不会等待前向传输完成:而是在模型结束为个别层附加 KV 缓存条目后立即启动 KV 页面复制。由于预填充请求可以分块,批处理调度器在执行前会通知 KV 信使当前计划的块。为了在跟踪层的同时支持 CUDA 图表,信使保持一个专用线程以在注意力输出投影后轮询增量计数器。该计数器仅在分片环境的主节点上保持:尽管 KV 缓存条目在附加后和注意力前有效,输出投影在层之间归约,隐式地同步它们。一旦观察到计数器变化,信使被通知并调用 fabric-lib 开始层的传输。

在最后一个块的传输完成后,任何额外的元数据也被复制过来:投机解码或 MTP 需要将对数值和隐藏状态移动到解码器。这些副本也通过 RDMA 执行,从预先分配的缓冲区中读取和写入。

在最后一块的所有挂起传输完成后,预填充节点释放 KV 页面并完成请求。解码节点并未明确通知:相反,它使用即时计数器跟踪已完成的操作数量。预填充端的 RDMA 操作数与传输的页面数成比例。当已知数量的页面和上下文副本完成时,fabric-lib 调用 KV 信使指示请求已准备好进行解码。信使释放任何上下文并将请求交给 LLM 引擎。

分片 KV 缓存传输

如果预填充和解码依赖于 Tensor 并行(TP)并相同地分片或复制 KV 缓存,则单个传输引擎协调多个设备以发送和接收所有副本的页面。为了能够在模型的执行器跨多个设备和进程复制的情况下使用单个信使和传输引擎,使用 cuMemcuMemImportFromShareableHandle 分配支持 KV 缓存的设备内存并将其映射到主进程中。传输引擎检查节点的拓扑以找到用于传输 KV 缓存切片的 NIC 和距离最近的 NUMA 节点的 CPU。

如果源和目标分片相同,传输是简单的,因为源和目标设备和页面之间存在一对一映射。在这种情况下,通过使用更多 GPU,可以使用更多相关的 NIC,达到更接近全带宽利用。如果存在不匹配关系,则传输引擎必须根据源和目标切片之间的比率拆分或重建页面。

如果预填充器将 KV 缓存拆分到更多设备上,解码器通过从预填充设备发送相应的一半来重建完整页面。如果解码器有更多的分片,它将接收来自多个源的页面。解码器需要了解预填充器的分片方案,以能够计算它预计将收到的 RDMA 写次数。如果涉及复制,预填充器会将设备分组到副本集中,这些副本集在其内部复制完整的 KV 缓存。目的地副本集被随机分配一个源集,以便使用所有可用设备发起 RDMA 写。

分片传输需要对 KV 缓存进行轻微调整。默认情况下,FlashInfer 依赖于 NHD 布局,它在头部内,页面中的令牌被排序。由于缓存可能按注意力头的数目进行分片,这会在头内产生不连续性。RDMA 传输不会隐式支持跨步写入,需要每个头执行一次传输操作。为了减少与 libfabric 的互动次数,我们通过采用 HND 布局组织 KV 缓存,将头维置于令牌数之前。这保证了连续性,允许通过一次写入复制整页。

投机性解码

投机性解码需要对分散的预填充-解码进行轻微调整。在我们的实现中,预填充节点不允许采样令牌。由于 Perplexity 的 Sonar 模型支持结构化输出,我们不希望在预填充器和解码器之间同步模式处理程序实现时遇到复杂性。在 MTP 和投机解码机制中,预填充草稿模型到最后一个令牌涉及从目标模型采样令牌。

为了绕过这些问题,预填充不包括输入序列的最后一个令牌。取而代之的是,在最后一个令牌之前的预填充隐藏状态或对数值被传输,并在解码器的下一步中被视为解码令牌。尽管这增加了延迟,因为在预填充后必须执行完整的解码步骤以发布第一个令牌,但实现的复杂性大大降低。

分散部署

我们已经部署或试验了不同型号的多种分散配置,以支持生产流量或内部评估工作负载。基于模型的大小和注意力机制,我们选择了适合的预填充器和解码器节点分片方案,以最佳利用 GPU。

DeepSeek-R1

在 DeepSeek 中,我们同时考虑了 Tensor-Parallel(TP)和 Data-Parallel(DP)部署。如在以前的博客文章中讨论的那样,TP 部署提供更好的延迟和降低吞吐量的代价,需要更多的 GPU 来服务繁重的流量。DP 部署在负载下扩展得更好,然而,由于设备间或节点间通信成本,峰值吞吐量较低。

DeepSeek 依赖于多头潜在注意力,压缩 KV 缓存。由于所有 KV 头都是压缩成一个单一的潜在向量,TP 不能对 KV 缓存进行分片,而必须在所有秩上复制潜在向量。分片发生在解压缩之后,因为每个秩可以从相同的潜在表示中提取不同的头。结果是,在预填充器和解码器分片中,所有 KV 缓存分片都是相同的。

在节点内 TP 设置中,预填充器和解码器都被相同的分片。传输从所有秩派遣,以充分利用所有可用的 NIC。然而,在 DP 部署中,TP 秩大小较小或每个 DP 秩分配给单个 GPU,任何预填充设备都持有 KV 缓存的复制副本可以派遣它。为了在所有可用 NIC 间平衡请求,我们随机选择一个 GPU 和一个 NIC,从预填充器到解码器发送 KV 缓存。

在混合预填充-解码中,由于频繁的预填充中断,R1 部署难以稳定地超过 50 TPS。然而,通过分离预填充,我们对每个请求触发了大约 100ms 的 TTFT 惩罚,但单个预填充节点可以在 3 个解码节点上保持一致的批量大小,在处理大约 1 QPS 的负载下(每个解码节点)提供超过 90 TPS 的吞吐量。对于数据并行的部署,TPS 略低,大约为 50,但实例可以处理每个秩 1 QPS 的负载,每个节点有 8 个秩。

Qwen3-Coder

此 480B 模型使用分组查询注意力(GQA),因此注意力可以轻松分片,并且可以在不牺牲 KV 缓存的内存下受益于张量并行。因此,我们可以在 8 个 GPU 上对模型进行预填充和解码分片,以便配对大约 3 个解码节点与一个预填充节点。由于注意力是分片的,我们依赖 HND KV 缓存布局来分片预填充器和解码器 KV 缓存,配对预填充秩与解码秩,并充分利用所有 NIC 并行传输切片。

有兴趣塑造我们API平台的未来吗?我们正在招聘。

加入我们的开发者社区,掌握最新版本、功能和更新。

有兴趣塑造我们API平台的未来吗?我们正在招聘。

加入我们的开发者社区,掌握最新版本、功能和更新。

有兴趣塑造我们API平台的未来吗?我们正在招聘。

加入我们的开发者社区,掌握最新版本、功能和更新。