分离预填充与解码
将预填充和解码分离到不同设备上可加速LLM推理并提高性能
为了从输入提示词中生成输出标记,LLM推理被分为两个阶段:预填充(prefill)和解码(decode)。预填充阶段在输入标记上运行,填充KV缓存,随后进入逐个生成标记的解码阶段。虽然单个解码步骤通常只需几十毫秒,但预填充所需的时间要长得多。如果两者在同一设备上运行,将预填充与解码混合会降低解码性能。在本文中,我们探讨了一种已建立的解决方案,即预填充与解码的分离(disaggregated),通过在不同设备上分别运行这两个阶段,从而最大化预填充吞吐量和解码延迟。
预填充与解码性能
在典型的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利用率,但缓存和矩阵乘法单元通常仍处于低利用状态。
当混合在一起时,包含预填充请求的批次会导致前向传递的延迟增加,从而对整个实例的解码吞吐量产生负面影响。虽然将预填充请求与解码请求混合或采用分块预填充(chunked prefill)可以略微改善解码性能,但很难保持足够的预填充吞吐量来在实例上处理足够的请求,以最大化解码吞吐量。对于具有典型输出长度的大型模型,为了保持较大的解码批处理规模,必须足够频繁地执行预填充,这会显著降低平均延迟并导致输出卡顿。
这些问题可以通过使用一组独立的节点来执行预填充和解码来解决。通过将一个预填充节点与多个解码节点关联,可以调度足够的请求进行预填充以最大化吞吐量,并保持解码节点上有足够多的并发请求以最大化解码吞吐量。预填充节点填充KV缓存,然后将其传输到解码节点。由于解码器不再需要为预填充而中断,延迟变得更加确定,因为活动请求不断增长的kv_len带来的总体影响显著减小。其代价是首字延迟(TTFT)的增加,因为通过网络传输KV缓存可能需要几十到几百毫秒。
KV Messenger
在Perplexity,我们关于分离预填充和解码的实现是围绕一个KV messenger构建的,它与LLM引擎交互,通过网络编排从预填充节点到解码节点的KV缓存传输。在预填充侧,messenger接收来自解码节点的请求,将其交给批处理调度程序,并跟踪前向传递执行,以便以尽可能低的延迟分发KV缓存。在解码侧,在分配了不可逐出页面后,messenger会阻止请求进入解码调度,直到收到KV缓存和解码器上下文传输完成的通知。
分离预填充需要高吞吐量、低延迟的连接,因此我们的实现针对RDMA进行了定制,支持EFA和ConnectX网络接口控制器(NIC)。KV Messenger构建在libfabric之上,利用我们的fabric-lib封装器在远程直接内存访问(RDMA)原语上提供更高级别的低延迟抽象,实现高效的页面和元数据传输,以及低延迟信号传输。在后台,fabric-lib协调GPU及其直接连接的NIC,将数据从预填充节点复制到解码节点。
收到请求后,预填充节点会分配相应的一组源KV页面,并使用其本地引擎调度预填充请求。为了最小化延迟,传输不会等待前向传递:相反,一旦模型完成将KV缓存条目附加到各层的KV缓存中,就会立即启动KV页面复制。由于预填充请求可以是分块的,批处理调度程序会在执行前通知KV messenger当前调度的块。为了支持CUDA图形同时能够跟踪层,messenger保留一个专门的问题,该问题在注意力输出投影后递增计数器并进行轮询。该计数器仅在分片环境中的主节点上维护:尽管KV缓存条目在附加后和注意力计算前有效,但输出投影会在各秩(rank)之间进行归约,从而隐式同步它们。一旦观察到计数器变化,messenger就会收到通知并调用fabric-lib来启动层的传输。
在最后一块的传输完成后,任何额外的元数据也会被复制过去:推测解码或MTP需要将logits和隐藏状态移动到解码器。这些复制操作也通过RDMA在预分配的缓冲区之间进行。
在最后一块的所有待处理传输完成后,预填充节点释放KV页面并完成请求。解码节点不会收到显式通知:相反,它使用即时计数器来跟踪已完成的操作数量。预填充侧的RDMA操作数量与传输的页面数量成正比。在完成已知数量的页面和上下文复制后,fabric-lib会调用KV messenger,指示请求已准备好进行解码。Messenger会释放所有上下文并将请求移交给LLM引擎。
分片KV缓存传输
如果预填充器和解码器依赖于张量并行(TP)并相同地分片或复制KV缓存,单个传输引擎会协调多个设备来发送和接收所有副本的页面。为了能够在模型执行器在多个设备和进程上复制的情况下使用单个messenger和传输引擎,cuMem和cuMemImportFromShareableHandle被用于分配支持KV缓存的设备内存,并将其映射到主进程中。传输引擎会检查节点的拓扑结构,以找到最接近NUMA节点的NIC和CPU,用于每个KV缓存切片的传输。
如果源和目标的分片方式相同,传输是非常简单的,因为源和目标的设备与页面之间存在一一映射。在这种情况下,分片隐式地帮助了传输延迟:通过使用更多的GPU,可以使用更多的关联NIC,从而更接近带宽利用率峰值。然而,如果存在不匹配,传输引擎必须根据源切片和目标切片之间的比例来拆分或重构页面。
如果预填充器将KV缓存拆分到更多设备上,则通过从预填充器设备发送相应的对半部分,在解码器上重构完整页面。如果解码器有更多分片,它会从多个源接收页面。解码器需要知道预填充器的分片方案,以便能够计算其预期接收的RDMA写入次数。如果涉及复制,预填充器会将设备分为副本集,在副本集内部复制完整的KV缓存。目标副本集会被随机分配一个源集,以便使用所有可用设备来发起RDMA写入。
分片传输需要对KV缓存进行轻微调整。默认情况下,FlashInfer依赖于NHD布局,该布局按头(head)内的标记顺序排列页面内的标记。由于缓存很可能是沿着注意力头数量进行分片的,这会在头内造成不连续性。RDMA传输不支持隐式的跨步写入,因此每次传输需要针对每个头执行一次操作。相反,为了减少与libfabric的交互次数,我们使用HND布局组织KV缓存,将头维度置于标记数量之前。这确保了连续性,允许通过单次写入操作复制页面。
推测解码
推测解码需要对分离的预填充-解码进行微调。在我们的实现中,不允许预填充节点采样标记。由于Perplexity的Sonar模型支持结构化输出,我们不希望承担在预填充器和解码器之间同步模式处理器实现的复杂性。在MTP和推测解码机制中,将草稿模型预填充到最后一个标记涉及从目标模型采样标记。
为了解决这些问题,预填充不包括输入序列的最后一个标记。相反,预填充在最后一个标记之前的隐藏状态或logits会被传输,并在解码器的下一步中将其视为解码标记。虽然这略微增加了延迟,因为必须在预填充后执行完整的解码步骤才能发出第一个标记,但这极大地降低了实现的复杂性。
分离部署
我们已经部署或试验了多种具有不同模型的分离配置,以支持生产流量或内部评估工作负载。根据模型的大小和注意力机制,我们为预填充和解码节点选择了合适的分片方案,以最大限度地利用GPU。
DeepSeek-R1
对于DeepSeek,我们考虑了张量并行(TP)和数据并行(DP)部署。正如之前博客文章中所讨论的,TP部署提供了更好的延迟,但代价是吞吐量较低,需要更多GPU来处理繁重的流量。DP部署在负载扩展方面表现更好,但由于设备间或节点间通信的开销,其峰值吞吐量较低。
DeepSeek依赖于多头潜在注意力(Multi-Head Latent Attention),压缩了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来并行传输切片。
引用
@online{perplexity2025DisaggregatedPrefillAndDecode,
author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
title = {{Disaggregated prefill and decode}},
date = {2025-08-01},
url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}