Pico-vLLM 开发日志 #12 PD分离

PD分离可以说是最近几年最为火热的llm推理引擎架构创新了。从Meta到Openai到Anthropic，内部几乎都肯定有完整的PD分离框架，其中Meta甚至因为这个发过相关的技术论文。它的核心卖点在于推理过程中的一个计算特性区分：Prefill是一个计算强度和prompt length相关的操作，是计算密集的，同时是单步重负载的；而Decode是一个计算强度几乎为固定值的操作，同时也是单步相对轻负载的。在前面的设计我们知道，基于Continuous Batching技术+Chunked Prefill会引入Prefill和Decode的混编，而这种混编会引起TTFT的抖动，以及其他方面的性能影响。如何解决它？答案意外的简单，把它们物理上分开不就行了？

但有意思的是，vllm直到最近还在把PD分离作为一个有experimental标签的“实验性”特征，25年8月才修好P侧和D侧的TP异构导致的不兼容问题。对于一个如此热门的开源项目来说，似乎进展相当慢了。所以在做之前的当时就已经在怀疑，是否可能有坑？做到目前来看的结果，确实是相对来说复杂度较高的feature，原因主要不是技术原理，而是工程实现的复杂度，尤其是和其他并行模式的耦合的复杂度。下面主要记录一下这部分产生的心得。

为什么需要PD分离

这个问题不难理解，在前面也讲到了，本质是计算强度的不同，导致的最优化的ridge point完全不同，进而导致硬件资源的非充分利用。不过在阅读论文之后，我发现其实还有一个原因，就是用户端的体验问题。下面将会逐一讲解。

TTFT和ITL的权衡

TTFT（Time To First Token）衡量的是从用户输入prompt开始，到模型开始生成第一个Decoding Token的时间间隔。这本质上是一个人机工程学指标：毕竟硬件的效率就摆在那里，TTFT是一个“权衡”，而不是算力上限的最精确衡量。但是它对于用户来说是非常重要的：我们在TTFT时间之后才可以真的开始阅读LLM生成的内容、进入非闲置的工作状态，在此之前是纯粹无效等待。它也很可能影响“系统是不是卡了”的直觉判断，一个TTFT高或者不稳定的LLM框架即使后续生成很快，也会无法高效维持用户的注意力。

ITL（Inter-Token Latency）是生成过程中相邻两个token之间的间隔。它决定的是"阅读体验是否流畅"。它同样会影响到阅读体验：如果大量Prefill突然加入，ITL就一定会升高（这是不可避免的），从而让阅读被打断。更重要的是人在阅读情况下对ITL比起TTFT更敏感，因为人没有被打断的心理预期。

这两者之间构成trade-off。如果为了保护ITL，限制每步最多只做一个Prefill（或者限制Prefill的token数），那TTFT就会变长，因为新请求要排队等。如果为了降低TTFT，激进地插入Prefill，那正在Decode的请求的ITL就会被打断。这是一个调度层面无法根治的矛盾，其问题在于“不稳定性”和“不可预测性”。只要Prefill和Decode在同一张卡上，这个跷跷板就存在。

在Continuous Batching框架里，Prefill和Decode共享同一个GPU。当一个新请求进来做Prefill的时候，它的计算量远大于同batch里正在Decode的那些请求（一个Prefill可能处理几百上千个token，而每个Decode只处理1个token）。这意味着这一步的总耗时被Prefill主导，所有正在Decode的请求都必须等这一步跑完才能拿到自己的下一个token。结果就是：正在生成中的请求的ITL被新请求的Prefill拖长了。确保它的稳定性在非分离情况下需要非常非常复杂的调度策略，几乎接近于时序预测而非确定性的调度策略，这就自然催生PD分离需求。

硬件资源的充分利用

更深层次的原因是Prefill和Decode的计算特征差异大到它们的最优硬件配置完全不同，放在一起跑意味着任何时刻都有资源在浪费。正如前面说的，它们具有不同的总体算术强度。我们都知道，硬件的最优化利用是当程序处于其ridge point，或者接近其ridge point的时候。但一个硬件的ridge point显然只有一个，而Prefill-Decode两个不同需求有两个在Roofline模型里差异极大的算术强度点（具体可以参加前面的博客）。当Prefill在跑的时候，HBM带宽没有被充分利用（compute-bound，Tensor Core和CUDA Core是瓶颈）。当Decode在跑的时候，GPU的算力没有被充分利用（bandwidth-bound，HBM是瓶颈）。这意味着无论在哪个阶段，都有一部分的硬件能力在闲置。

PD分离带来一个潜在机会。如果我们进行恰当的分离，是否可以部署异构硬件，从而以相同的预算成本做相同的事情？Prefill用算力更强的卡，Decode用显存带宽更高或者成本更低的卡。实际上，最近几年Nvidia和华为也都在基于Prefill和Decode阶段的不同计算特性，设计具有不同ridge point的GPGPU（Rubin CPX）或者NPU硬件变体。它们基本上也就是为了匹配PD分离里分离不同算术强度的操作而设计的，避免资源闲置和同成本下的利用效率最大化。

即使是在软件层面，也已经有了（实际上更完善的）软件PD分离支持，这部分可以参考NVIDIA Dynamo（官方的PD分离调度框架）和NIXL（Inference Xfer Library），在此不过度展开了。

可用的工具箱

对于具体的PD分离实现，其实有很多种不同选择，以下几种均有实现的潜力。为了单人开发的复杂性和可维护性控制，笔者选择的是直接使用torch.distribute。然而，此处也列出工业界和其他可用的选项，供大家参考。这部分很大程度上是通过ai提供思路，然后再进行调研的，再次感谢ai。

使用torch.distribute（NCCL）

这是最容易上手的方案，也是之前Pico-vLLM在TP的开发阶段已经用过的工具。torch.distributed封装了NCCL的集合通信原语，提供了send()/recv()的点对点通信接口。

优点：开发成本低。NCCL对NVLink、PCIe、InfiniBand的底层拓扑发现和路径选择都是自动的，不需要开发者手动管理。

缺点：NCCL的设计初衷是集合通信，而不是点对点的大块数据传输。另一个问题是NCCL要求通信的所有进程在初始化时就加入同一个通信组（process group），这意味着Prefill Worker和Decode Worker必须在启动时就互相知道对方的存在。这在静态部署下没问题，但如果想要做动态的弹性扩缩容（比如根据负载动态增减Prefill/Decode Worker的数量），NCCL的这种静态组模型就会变得非常不灵活。此外，如果想在这个方案下做，需要先把传输内容gather到一起，进行传输，然后再在远端scatter。这意味着两次额外拷贝。

使用NIXL

NIXL是NVIDIA在GTC 2025上开源的、专门为推理场景设计的点对点数据传输库。它是NVIDIA Dynamo推理框架的核心数据搬运层。和NCCL不同，NIXL从设计之初就是为了解决推理场景下的KV cache传输这个具体问题的。

NIXL的核心抽象是Agent和可插拔的传输后端。每个Worker（无论是Prefill还是Decode）运行一个NIXL Agent。Agent向NIXL注册自己的内存区域（GPU HBM、CPU DRAM、NVMe SSD等），然后通过异步的transfer()接口发起传输。NIXL会自动选择最优的传输后端。

优点：专为推理KV cache传输设计；支持异步、非阻塞传输；支持非连续内存scatter/gather；多后端自动选择；与vLLM、SGLang、TensorRT-LLM等主流框架已有集成。

缺点：依赖较重（需要UCX、RDMA驱动等），在没有RDMA的环境下（比如我们的autodl PCIe环境）退化为TCP。NIXL目前仍在快速迭代中，API稳定性和文档完善度还有提升空间。安装也不是特别顺畅。此外更重要的是，它是C原生的，将其整合进入Python需要额外的大量努力。

使用CUDA IPC

CUDA IPC是CUDA原生提供的跨进程GPU内存共享机制。核心API是cudaIpcGetMemHandle()和cudaIpcOpenMemHandle()：前者从一块已分配的GPU内存中导出一个可序列化的handle，后者在另一个进程中用这个handle获取一个指向同一块物理显存的指针。

优点：同机部署下理论上是延迟最低的方案。因为它是零拷贝的，传输延迟为零。实现也相对直接，不需要额外的库依赖。

缺点：只能在同一台机器上使用。这是因为它需要同一块物理显存。如果要通过网络，自然就没有所谓的共享虚拟内存空间了。

设计模式

基于实例角色role的兼容式设计

这部分的想法很简单：在不额外增加函数的情况下实现不同实例角色的指派，而在实现上保持兼容。

实例具有三种可能的角色：“p”、“d”和“pd”。

传输层抽象分离

一种简单的想法是把传输实现在Engine里。它是最简单的，但思考之后我并没有采纳这样的方案。原因是传输的具体方案实在是太多了：就像刚刚说的那样，不同的后端实现、同步异步、和Scheduler策略的耦合、P侧和D侧的收发约定将会带来不可维护的复杂性。因此不如单独抽象一个基类，然后作为接口来开发。这是软件工程的想法，但值得一提。目前实现的是同步模式，但是，传输层的分离将允许未来进行异步和更加复杂的可扩展性。

注意事项和Bug

值得注意的Bug汇总

这里总结了几个遇到的bug，详细记录一下，以免后面忘记再踩重复的坑，同时也给后来的读者一个提示，避免有相同的错误。

1、Python的可变默认值问题

在开发过程中，遇到类似下面的bug。当提交若干个（超过1个prompt的时候），出现了具有奇怪但是固定模式的输出乱码。如下图：

 [Request 1] The capital of France is Paris Paris 11200
AAExample   112

  [Request 2] 1 + 1 = Paris Paris 11200
AAExample   112

输入的prompt分别是：

1
2
3
4
5
prompts = [
    "The capital of France is",
    "The capital of France is",
    "1 + 1 =",
]

其中第一个request的结果被彻底吞掉，而。这个我一开始以为是KV cache相关的部分又出了问题，因为太像了。经过长达三个小时的详细排查，最终发现了一个bug的特性：它重复的Paris这个首个Decode token的数量，和input的prompt完全相关。当输入的prompt变成：

1
2
3
4
5
6
prompts = [
    "The capital of France is",
    "The capital of France is",
    "The capital of France is",
    "1 + 1 =",
]

的时候，错误的输出结果变成了：

[Request 3] 1 + 1 = Paris Paris Paris 1112000
AAAExample

又多了一个Paris。这给了我缩小范围的想法：它很可能是很多个prompt的Decode结果，同时挂在了所有的prompt的generated_ids序列上。那么，它的工作机制是什么呢？先看原始request类的代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class RequestStatus(Enum):
    WAITING  = "waiting"   # 在 waiting 队列，还没做 prefill
    PREFILL  = "prefill"   # 本步正在做 prefill
    DECODING = "decoding"  # 已经 prefill 完，正在 decode
    FINISHED = "finished"  # 已完成

''' 请求对象，包含请求的所有信息和状态
- request_id: 请求 ID，唯一标识一个请求
- input_ids: 输入的 token ids，shape (1, init_seq_len)，包含整个 prompt
- generated_ids: 已经生成的 token ids，shape (1, )，初始为 空
- max_new_tokens: 最多生成多少个 token
- temperature: 采样温度，传递给 sampler
- top_p: top-p 截断，传递给 sampler
- kv_cache: 每个请求独享一个 KV cache 实例，存储生成过程中的 KV 状态
'''
class Request:
    request_id: int
    input_ids: List[int]  # (1, init_seq_len)，包含整个 prompt
    generated_ids: List[int]  # (1, )，包含已经生成的 token ids，初始为 空
    max_new_tokens: int
    temperature: float
    top_p: float
    kv_cache: PagedKVCache  # 每个请求独享一个 KV cache 实例
    request_status: RequestStatus
    has_eos_token: bool  # 是否已经生成 eos_token，scheduler 不直接接触 tokenizer 和 eos_token_id，这个由 engine 在 decode_step 后更新

    def __init__(self, request_id: int, input_ids: List[int], max_new_tokens: int, temperature: float, top_p: float, kv_cache: PagedKVCache, generated_ids: List[int]=[]):
        self.request_id = request_id
        self.input_ids = input_ids
        self.generated_ids = generated_ids
        self.request_status = RequestStatus.WAITING
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.kv_cache = kv_cache
        self.has_finished_notification = False # engine改变这个状态，scheduler根据这个状态改变 request_status和移出队列

    def is_max_len_finished(self) -> bool:
        return len(self.generated_ids) >= self.max_new_tokens

    @property
    def prompt_len(self) -> int:
        return len(self.input_ids)
    
    @property  
    def total_len(self) -> int:
        return len(self.input_ids) + len(self.generated_ids)

观察这一句：

1
def __init__(self, generated_ids: List[int] = []):

它涉及到一个Python语法：Python的函数默认值在函数定义时求值一次，不是每次调用时。[] 在 Python 解释器加载这个类定义的时候就创建了一个 list 对象，存在 Request.__init__.__defaults__ 里。之后每次调用 Request(...) 不传 generated_ids 时，用的都是同一个 list 对象。所以 r0.generated_ids.append(42) 之后，r1.generated_ids 也变成了 [42]，因为它们指向内存中同一个 list。这是一个非常隐蔽的bug，因为大部分测试原本是在B=1情况下验证，此时只有一个Request实例，共享不共享都无所谓。B>1时多个 Request 往同一个list里append，所有请求的 generated_ids[-1] 返回同一个值，从外部看起来，就是decode出了相同的token。

在修复了Request类中相关的错误代码之后，这个问题就成功解决了，框架开始输出正确的request结果。如图：

  [Request 0] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
  [Request 1] The capital of France is Paris. The capital of the United States is Washington, D.C. The capital of Canada is Ottawa
  [Request 2] 1 + 1 = 2
Example 2:
Input: nums = [1,2,3,4,

2、锁死问题

这个严格来说不是实现的bug问题，而是设计模式的问题。当我实现到一半，在写测试脚本的时候，我意识到一个问题：对于一个有并发的项目，接受方怎么确认不用再进行通信了？这个问题乍一看可以通过预先约定次数来解决，但实际上比这个复杂一些，因为很多情况下你根本没有办法提前预知次数，当然也就没法约定它。又考虑了，如果传递一个空请求，那么就自动结束，不再接收信息。这个逻辑在同步下有效，但是在异步下无效：你如何确认在Prefill端没有请求之后，就不会再出现请求了？也许只是短暂出现了波谷，而不是真的停止了。但如果这样不退出的话，那么实际上你没有办法在Engine内部设置一个正常的退出条件，只能通过狼狈的ctrl+c来强行终止。思来想去，最后的结论是：无法在Engine内部做出这个判断，它应该是用户层控制的部分。

为此进行了重构，现在，Engine在用完之后，用户需要显式调用一次“no more requests”函数，以此显式通知Prefill侧的Engine不再接受新request的submit，并且在清空流水线之后通知Decode侧，让它在清空流水线之后同样退出。这样，就比较好的解决了锁死问题。