P2P NCCL 连接器¶
一种基于点对点通信并支持动态伸缩的 xPyD 实现,部分灵感来源于 Dynamo。
详细设计¶
整体流程¶
如图 1 所示,该 PD 解耦 (PD disaggregation) 方案的整体流程通过一个请求流来描述:
- 客户端向代理/路由 (Proxy/Router) 的
/v1/completions接口发送 HTTP 请求。 - 代理/路由通过轮询或随机选择算法选中一个 1P1D (1 个 Prefill 实例 + 1 个 Decode 实例),生成一个
request_id(规则稍后介绍),将 HTTP 请求中的max_tokens修改为 1,然后将请求转发给 P 实例。 - 紧接着,代理/路由将 原始 HTTP 请求 转发给 D 实例。
- P 实例 执行 Prefill (预填充),然后 主动将生成的 KV Cache 发送 给 D 实例(使用 PUT_ASYNC 模式)。D 实例的
zmq_addr可以通过request_id解析得到。 - D 实例 拥有一个 专用线程 用于接收 KV Cache(以避免阻塞主进程)。接收到的 KV Cache 被保存到 GPU 显存缓冲区 中,其大小由 vLLM 启动参数
kv_buffer_size决定。当 GPU 缓冲区满时,KV Cache 将存储在 本地 Tensor 内存池 中。 - 在 Decode (解码) 过程中,D 实例的主进程从 GPU 缓冲区 或 内存池 中获取 P 实例传输的 KV Cache,从而 跳过 Prefill 阶段。
- 完成 Decode 后,D 实例将结果返回给 代理/路由,再由其转发给 客户端。
代理/路由 (演示)¶
一个简单的 HTTP 服务作为客户端请求的入口,并启动一个后台线程监听 P/D 实例上报的 HTTP IP 和端口,以及 ZMQ IP 和端口。它维护一个 http_addr -> zmq_addr 的字典。http_addr 是 vLLM 实例处理请求的 IP:PORT,而 zmq_addr 是用于 KV Cache 握手和元数据接收的地址。
代理/路由负责根据客户端请求的特征(如 Prompt)选择 1P1D 并生成相应的 request_id,例如:
目前为了快速验证 xPyD 的可行性,采用了 1P1D 的轮询选择策略。未来计划结合 Trie 树和实例的负载状态来选择合适的 P 和 D。
每个 P/D 实例定期向代理/路由发送心跳包(当前为每 3 秒一次)进行注册(即上报 http_addr -> zmq_addr)并保持连接。如果某个实例崩溃并长时间未能发送心跳,代理/路由将移除该超时的实例(此功能尚未开发)。
KV Cache 传输方法¶
KVCache 传输有三种方法:PUT、GET 和 PUT_ASYNC。这些方法可以通过 --kv-transfer-config 和 kv_connector_extra_config 参数(通过 send_type 字段)指定。PUT 和 PUT_ASYNC 均涉及 P 实例主动向 D 实例发送 KVCache。区别在于,PUT 是一种阻塞主进程的同步传输方式,而 PUT_ASYNC 是一种异步传输方式。PUT_ASYNC 使用专用线程发送 KVCache,因此不会阻塞主进程。相比之下,GET 方法涉及 P 实例在完成 Prefill 计算后将 KVCache 保存到内存缓冲区。D 实例在分配好空间后,主动从 P 实例获取计算好的 KVCache。
实验结果表明,这些方法的性能从高到低依次为:PUT_ASYNC → GET → PUT。
基于 ZMQ 和 NCCL 的 P2P 通信¶
只要知道对方地址,就可以进行点对点 KV Cache 传输(使用 NCCL),且不受 Rank 和 World Size 的限制。这支持了 PD 解耦场景下的动态伸缩。这意味着增加或删除 P/D 实例不需要重启整个系统。
每个 P/D 实例只需创建一个 P2pNcclEngine 实例。该实例维护一个 ZMQ Server,运行专用线程监听 zmq_addr 地址,并接收来自其他实例的控制流请求。这些请求包括建立 NCCL 连接的请求和发送 KVCache 元数据(如张量形状和数据类型)的请求,但并不实际传输 KVCache 数据本身。
当 P 实例和 D 实例首次传输 KVCache 时,它们需要建立 ZMQ 连接和 NCCL 组。对于后续的 KVCache 传输,这些连接和组将被复用。NCCL 组仅包含两个 Rank,即 World Size 等于 2。此设计旨在支持动态伸缩,只要已知对方地址,即可进行点对点 KVCache 传输,而不受 Rank 或 World Size 的限制。
NCCL 组拓扑¶
目前仅支持对称 TP(张量并行)方式的 KVCache 传输,未来将支持非对称 TP 和 PP(流水线并行)。图 2 展示了 1P2D 设置,其中每个实例的 TP 度为 2。总共有 7 个 NCCL 组:三个 vLLM 实例各自拥有一个 TP=2 的 NCCL 组。此外,P 实例的第 0 号 GPU 卡与每个 D 实例的第 0 号 GPU 卡建立一个 NCCL 组。同理,P 实例的第 1 号 GPU 卡与每个 D 实例的第 1 号 GPU 卡建立一个 NCCL 组。
每个 NCCL 组会占用一定的 GPU 显存缓冲区用于通信,大小主要受 NCCL_MAX_NCHANNELS 环境变量影响。当 NCCL_MAX_NCHANNELS=16 时,一个 NCCL 组通常占用 100MB;当 NCCL_MAX_NCHANNELS=8 时,占用 52MB。对于大规模 xPyD 配置(如 DeepSeek 的 96P144D),此实现目前不可行。展望未来,我们正考虑使用 RDMA 进行点对点通信,并持续关注 UCCL。
GPU 显存缓冲区与 Tensor 内存池¶
关于显存缓冲区的权衡如下:对于 P 实例,PUT 和 PUT_ASYNC 模式下不需要显存缓冲区,但 GET 模式下是必需的。对于 D 实例,所有三种模式都需要显存缓冲区。D 实例的显存缓冲区不宜过大;同样,对于 GET 模式下的 P 实例,缓冲区也不宜过大。D 实例的显存缓冲区用于暂存 P 实例发送的 KVCache。如果过大,会减少 D 实例正常推理可用的 KVCache 空间,从而降低推理的 Batch Size,最终导致输出吞吐量下降。显存缓冲区的大小由 kv_buffer_size 参数配置(单位为字节),通常设置为显存大小的 5%~10%。
如果 P 实例的 --max-num-seqs 参数设置较大,由于 Batch Size 较大,P 实例会同时生成大量的 KVCache。这可能超过 D 实例显存缓冲区的容量,导致 KVCache 丢失。一旦 KVCache 丢失,D 实例就需要重新计算 Prefill,相当于进行了两次 Prefill。因此,首字延迟 (TTFT) 将显著增加,导致性能下降。
为了解决上述问题,我设计并开发了一个用于存储 KVCache 的本地 Tensor 内存池,灵感来源于 Linux 内存模块中的伙伴系统 (Buddy System)。由于服务器内存足够大(通常在 TB 级别),无需考虑前缀缓存或基于块的设计来重用内存,从而节省了空间。当显存缓冲区不足时,KVCache 可以直接存储在 Tensor 内存池中,D 实例随后从内存池中检索。读写速度即 PCIe 速度,PCIe 4.0 的速度约为 21 GB/s,通常快于 Prefill 速度。否则,Mooncake 和 lmcache 等解决方案也就没有必要了。Tensor 内存池充当了防洪蓄水池,通常仅在突发流量时使用。在最坏的情况下,我的解决方案也不会比现有的 Cache 存储方案表现更差。
安装 vLLM¶
运行 xPyD¶
使用说明¶
- 以下示例在 A800 (80GB) 设备上运行,使用 Meta-Llama-3.1-8B-Instruct 模型。
- 请注意
kv_buffer_size(单位:字节)的设置。经验值为 GPU 显存大小的 10%。这与 kvcache 大小有关。如果设置过小,用于暂存接收到的 kvcache 的 GPU 显存缓冲区会溢出,导致 kvcache 被存储到 Tensor 内存池中,从而增加延迟。如果过大,用于推理的 kvcache 空间会减少,导致 Batch Size 减小,吞吐量下降。 - 对于 Prefill 实例,在非 GET 模式下,
kv_buffer_size可以设置为 1,因为 Prefill 当前不需要接收 kvcache。但在 GET 模式下,需要较大的kv_buffer_size,因为它需要存储发送给 D 实例的 kvcache。 - 如果发生冲突,您可能需要在以下命令中修改
kv_buffer_size和port。 PUT_ASYNC提供最佳性能,应优先选用。--port必须与--kv-transfer-config中的http_port一致。disagg_proxy_p2p_nccl_xpyd.py脚本将使用端口 10001(用于接收客户端请求)和端口 30001(用于接收来自 P 和 D 实例的服务发现请求)。- 运行代理的节点必须安装
quart。 - 支持多节点;只需修改
--kv-transfer-config中的proxy_ip和proxy_port即可。 - 在以下示例中,假设 代理的 IP 为 10.0.1.1。
运行 1P3D¶
代理 (例如 10.0.1.1)¶
cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
python3 disagg_proxy_p2p_nccl_xpyd.py &
Prefill1 (例如 10.0.1.2 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20001 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
Decode1 (例如 10.0.1.3 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20002 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
Decode2 (例如 10.0.1.4 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20003 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
Decode3 (例如 10.0.1.5 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20004 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
运行 3P1D¶
代理 (例如 10.0.1.1)¶
cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
python3 disagg_proxy_p2p_nccl_xpyd.py &
Prefill1 (例如 10.0.1.2 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20001 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
Prefill2 (例如 10.0.1.3 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20002 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
Prefill3 (例如 10.0.1.4 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20003 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
Decode1 (例如 10.0.1.5 或 10.0.1.1)¶
命令
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20004 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
单次请求¶
curl -X POST -s http://10.0.1.1:10001/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "base_model",
"prompt": "San Francisco is a",
"max_tokens": 10,
"temperature": 0
}'
基准测试¶
命令
vllm bench serve \
--backend vllm \
--model base_model \
--tokenizer meta-llama/Llama-3.1-8B-Instruct \
--dataset-name "random" \
--host 10.0.1.1 \
--port 10001 \
--random-input-len 1024 \
--random-output-len 1024 \
--ignore-eos \
--burstiness 100 \
--percentile-metrics "ttft,tpot,itl,e2el" \
--metric-percentiles "90,95,99" \
--seed $(date +%s) \
--trust-remote-code \
--request-rate 3 \
--num-prompts 1000