这次迭代完成的是两个feature：异步模式的PD分离实现，以及和异构并行度下的PD分离-TP耦合的设计兼容。这部分做完之后，就只剩下基于Radix-Tree的Prefix Caching特性需要实现了，随后就进入整理、代码优化和交互界面设计阶段。这篇博客主要记录一下同步-异步模式两个的思路和实现，包括生产者-消费者模型；以及基于异构的PD分离-TP并行设计的具体实现。

什么是同步和异步模式

所谓同步模式，指的是"存在一个可序列化的顺序"，即整个请求的生命周期可以被描述为一个严格的线性序列：Prefill → KV Cache传输 → Decode。换言之，每一步不会在完成之前开始下一步，控制流是单一的。

而所谓异步模式，指的是"不存在这样一个全局的串行顺序"，各阶段之间通过显式的信号机制进行协调，而非通过阻塞等待来保证先后关系。Prefill端完成KV Cache的生产后，不需要等待Decode端确认接收就可以立刻开始处理下一个请求；传输过程可以和两端的计算并发进行；Decode端在KV Cache就绪后被通知并开始工作。控制流不再是单一的，而是分裂为多条独立推进的流，它们之间只在必要的数据依赖点上同步。

这两者的区别在于，当每一步不存在严格的对上一步的依赖，且对硬件资源的依赖不相互重叠的时候，即可通过这种方式进行某种意义上的重叠，从而提高性能。否则，两者的区别就不是很大，异步可能因为信号的轮询反而有性能的额外开销。

同步模式

同步模式的设计很简单。它的流程是：请求到达Prefill端，完整执行所有层的Prefill，生成全部层的KV Cache，然后一次性将KV Cache传输到Decode端，Decode端接收完毕后将请求加入Decode调度队列，开始逐token生成。它当然很容易保证和调试正确性，但它毫无疑问对于生产级来说是不可取的。

异步模式

异步模式的设计复杂度则比预想中高很多。它的最主要障碍反而不在于具体的编码实现，而在于设计模式的收发约定：未来还有多少个请求一定会要发送（发送方和接收方不再同步执行，那接收方怎么知道"该准备接收了"，怎么知道“不用再接收了”），以及后端的特性支持情况：保序性、是否支持tag以区分信道。不幸的是，对于这次实现的torch.distribute使用的nccl后端来说，它保序但不支持tag区分。而对于前者，实际上这没有可以简单解决的方法。

具体来说，NCCL的P2P是双边操作：isend必须和对端的irecv配对才能真正开始传输，单方面调用isend并不会把数据推出去。这意味着Decode端必须提前挂好irecv等着，但问题是它不知道Prefill端什么时候会发、甚至不知道还有没有下一个请求要发。

一个请求的KV Cache传输至少需要三次P2P通信：size、meta、data。在不知道后续是否还有请求的情况下，Decode端无法为下一个请求提前挂出irecv。这直接导致了一个硬约束：KV传输在请求级别是串行的，同一时刻只能有一个请求的传输流程在进行。 上一个请求的三次握手没有走完，下一个请求的传输就无法开始。

当然，更复杂的设计其实是存在的：比如在消息体中携带后续请求的元信息，或者建立一条独立的控制信道来提前广播传输计划，从而让Decode端可以预先挂出多个irecv实现传输的流水线化。但这已经不是推理引擎层面的工作了，而是进入了高性能通信协议栈的设计范畴，超出了本项目的边界。

NIXL的必要性

前面说了这么多，其实，目前的NCCL的局限已经体现的很明显了。这就自然引出了NIXL的需求：如果我们不需要双边的P2P语义，这个问题就根本不存在了。让ai老师介绍一下它的特长所在：

**NIXL（NVIDIA Inference Xfer Library）**是NVIDIA在GTC 2025上开源的、专门为分布式推理场景下的数据搬运设计的点对点传输库。它和NCCL的本质区别在于通信语义：NCCL的P2P是双边的send/recv模型，发送方和接收方必须配对才能完成传输；而NIXL采用的是单边的read/write模型——一个节点通过带外网络（如TCP或ETCD）导出自己的KV Cache元数据，其他节点拿到元数据后可以直接读写远端的GPU显存，远端甚至不需要在那一刻做任何配合操作。这意味着前面讨论的双边配对约束在NIXL的模型下天然不存在：Decode端可以在任意时刻、对任意数量的请求发起KV Cache读取，不需要和Prefill端一一握手。

在底层传输机制上，NIXL支持RDMA（InfiniBand和RoCE）、GPU-Direct RDMA、GPUDirect Storage、NVMe-oF、POSIX socket、以及S3兼容的对象存储等多种后端，并且会根据硬件配置自动选择最优的传输路径。一个关键的性能特性是，NIXL的传输不消耗GPU的SM资源——数据通过GPU-Direct RDMA直接在网卡和显存之间搬运，不经过GPU的计算单元，因此不会和模型的forward计算争抢算力。在实际benchmark中，NIXL在256KB到1MB的典型KV Cache传输大小上比NCCL快30%到50%。

NIXL的核心实现是C++，通过pybind11提供Python绑定，同时也提供Rust绑定。它目前已经被集成到NVIDIA Dynamo、TensorRT-LLM、vLLM、SGLang、LMCache等主流推理框架中，是NVIDIA推理基础设施栈中负责"搬数据"这一层的标准组件。

在后续实现中，也有计划在实现所有核心feature之后，将其变为PD分离的数据搬运的新后端，并且进行benchmark评测。

生产者-消费者模型

生产者-消费者模型是并发编程中最经典的协作模式之一。它描述的是这样一种场景：一方负责生产数据，另一方负责消费数据，两者通过一个共享的缓冲区解耦，各自以独立的节奏运行。生产者不需要等消费者处理完上一份数据才能继续生产，消费者也不需要等生产者准备好下一份数据才能开始处理。只要缓冲区里有东西，消费者就取；只要缓冲区没满，生产者就放。两者之间唯一的同步点是缓冲区的状态：空的时候消费者等待，满的时候生产者等待。

这个模型对通信机制的要求也很明确：第一，需要一个共享的、线程安全的缓冲区；第二，需要一种通知机制让双方知道缓冲区的状态变化（通常是信号量、条件变量或事件）；第三，生产者写入缓冲区的操作和消费者从缓冲区读取的操作不应该要求对方的主动配合。换句话说，理想的通信模式是单边的，生产者侧和消费者侧不需要握手，也不需要提前约定通信的固定总量。

PD分离就是典型的生产者-消费者场景。Prefill端生产KV Cache，Decode端消费KV Cache，两者的处理速度天然不同（Prefill是compute-bound的大块计算，Decode是memory-bound的逐token生成），需要解耦才能各自高效运行。但从实现上看，能否真正实现生产者-消费者模型，取决于底层通信机制是否满足上述要求，特别是第三点：单边性。如果传输机制要求双边配对（如NCCL的isend/irecv），那生产者每次放数据都需要消费者同时伸手来接，这就把生产者-消费者退化成了一个双方必须同步握手的模型，缓冲区的解耦作用被大幅削弱。而如果传输机制支持单边读写（如RDMA），生产者只需要把数据写入一块预先注册的远端内存，消费者在任意时刻来读就行，这是生产者-消费者模型真正能发挥作用的前提。

基于异构的PD分离-TP并行设计

这部分的思想同样也很简单，就是让P侧占用的卡数量（并行度）可以和D侧不同，两者解耦。这么做的好处显而易见如果P和D的计算特性差异很大，那么，给两者分配一模一样的GPU数量很可能是不明智的。此外，它也会影响在特殊集群或者有掉卡情况下集群的弹性扩缩容。而如果P侧和D侧的并行度可以不同，那么整个性能调整就可以很自由，可以通过精细的调优最大化集群硬件的利用率。

KV Cache的重分片

KV Cache跨TP重分片的实现是核心的开发需求。问题的根源在于：TP并行下，每张卡只持有KV heads的一个子集。以Qwen2.5-1.5B为例，模型有2个KV heads，如果Prefill端TP=2，每张卡各持有1个KV head的Cache；而如果Decode端TP=1，单卡需要完整的2个KV heads。这意味着KV Cache的传输不再是简单的整块搬运，而是必须在传输过程中完成一次重新拼合或拆分——从Prefill端的多张卡上各取一部分，聚合成Decode端所需要的完整形状，或者反过来将完整的Cache切分后分发到Decode端的多张卡上。

设模型共有 $N_{kv}$ 个KV heads。在TP并行度为 $tp$ 的情况下，rank $r$ 持有的KV heads范围是：

因此，Prefill端rank $r_p$（并行度 $tp_p$）持有的heads范围是：

Decode端rank $r_d$（并行度 $tp_d$）需要的heads范围是：

重分片的本质就是计算 $H_p(r_p)$ 和 $H_d(r_d)$ 之间的交集。对于Decode端的某个rank $r_d$，它需要从所有满足 $H_p(r_p) \cap H_d(r_d) \neq \emptyset$ 的Prefill端rank收集数据。具体来说，$r_d$ 需要接收数据的Prefill端rank范围是：

而从每个这样的 $r_p$ 处，$r_d$ 需要取出的head局部索引范围（相对于 $r_p$ 自己持有的heads的起始位置）是：

当 $tp_p = tp_d$ 时，所有公式退化为 $r_p = r_d$ 的一对一映射，不需要任何重分片——这也是同构PD分离的特殊情况。

写成公式就是这么几行，但翻译到代码里就会发现：每个传输操作的source rank、目标rank、在KV Cache tensor上的slice范围、以及block_table中的物理块重映射，全都要从这些公式算出来，而且必须在每种 $tp_p$、$tp_d$ 的组合下都正确。这还只是考虑了p、d互相的差异是整数倍的情况。如果不是整数倍、或者并行度大于KVhead的数量，其复杂度将会加倍。非整数倍意味着某个Decode端rank需要的heads会横跨某个Prefill端rank持有的heads区间的中间位置，产生不对齐的切片；而并行度大于KV head数量则需要引入GQA的head重复映射，同一个物理KV head的Cache要被复制到多张卡上。在pico-vllm的实现中，默认p和d的分配一定是整数倍的。

此外有意思的是，vLLM的异构并行度PD分离支持长期处于实验状态，相关的bug和兼容性问题从2024年底的路线图规划到2025年下半年仍在持续修复中。也许从中可以一窥跨并行度KV重映射在边缘情况下的实现难度有多令人头疼。