vLLM的PD分离

vLLM的PD分离

vLLM的PD分离是指vLLM的Prefill和Decode分离到不同的实例中执行。

新增配置

新增 KVTransferConfig 配置，决定了实例的类型。如果是 prefill 则 role 为 producer，如果是 decode 则 role 为 consumer，并且要设定传输的方法。

• is_kv_transfer_instance 判断是否是 PD 分离的实例。

代码实现

在 ./vllm/worker/model_runner.py 中：

1. 在计算之前执行 need_rev_kv，检查是否是 consumer，且当前 run 是不是 prefill。然后调用 get_kv_transfer_group().recv_kv_caches_and_hidden_states。
2. 在计算之后执行 need_send_kv，检查是否开启配置 producer，并且当前 run 是不是 prefill（对于以前未分离的结构来说，decode实例要经历prefill阶段，
   但是prefill已经被prefill实例做掉了，所以要等着接受prefill的KVCache，不需要重复计算prefill了）。然后调用 get_kv_transfer_group().send_kv_caches_and_hidden_states。

KVTransfer 实例

get_kv_transfer_group 会返回一个 KVTransfer 的实例，是一个全局实例，初始化方式如下。其中的 rank 0 代表 prefill，rank 1 代表 decode。

_KV_TRANSFER = kv_transfer.KVTransferAgent(
    rank=get_world_group().rank,
    local_rank=get_world_group().local_rank,
    config=vllm_config)

Transfer 实现

Transfer 的实现在 vllm/distributed/kv_transfer，这种解耦的设计是为了对接多种实现，比如 Mooncake 的开源 TransferEngine。Transfer 内部会调用 connector 的 send 和 recv 方法，这个方法是一个抽象方法，需要子类实现。目前有两种实现：Mooncake 的 transfer 和 PyNccl 的 transfer。

# Register various connectors here.
# The registration should not be done in each individual file, as we want to
# only load the files corresponding to the current connector.
KVConnectorFactory.register_connector(
    "PyNcclConnector",
    "vllm.distributed.kv_transfer.kv_connector.simple_connector",
    "SimpleConnector")

KVConnectorFactory.register_connector(
    "MooncakeConnector",
    "vllm.distributed.kv_transfer.kv_connector.simple_connector",
    "SimpleConnector")

Connector 依赖

Connector 依赖 kv_pipe 的实现。

• from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator 用来实现 PyNccl 的 kv_pipe。其中的 Send 和 Recv 会依赖 NCCL 的集合通信实现。
• 如果是 Mooncake pipe，import mooncake_vllm_adaptor as mva 这个模块，基于 ZeroMQ 的通信，通过 pickle 去序列化 tensor。

def _send_impl(self, tensor: torch.Tensor) -> None:
    """Implement the tensor sending logic."""
    value_bytes = pickle.dumps(tensor)
    self.transfer_engine.send_bytes(value_bytes)

KV Lookup Buffer 实现

另外还有一种 kv_lookup_buffer 的实现，抽象的接口是非阻塞 insert 和阻塞的 drop_select。

• Producer 调用 insert，consumer 调用 drop_select。目前 SimpleBuffer 也是基于 Pipe 去实现的，insert 变 send，drop_select 变 recv。
• 如果有一些中心化的 KVCacheBuffer 的话可能可以不用基于 Pipe 的实现。比如可以基于分布式的 LMCache?

Prefill 启动

vLLM 目前的实现是基于 connector 的。Prefill 的启动时通过设置 max_token 为 1 来执行，当生成了 bonus token 以后转而去调用 decode 的实例。