A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 01:00 UTC
[V1][P/D]Enhance Performance and code readability for P2pNcclConnector (#20906)
Signed-off-by: Abatom <abzhonghua@gmail.com>
Zhonghua Deng committed
Jul 17, 2025 at 13:13 UTC
8a4e5c5f3c1d39e924e48a87c9cc6cf382aa3532
4 files changed
+266
-256
docs/design/v1/p2p_nccl_connector.md
+30
-62
@@ -31,7 +31,7 @@ Each P/D instance periodically sends a heartbeat packet to the Proxy/Router (cur
31
32
## KV Cache Transfer Methods
33
34
-There are three methods for KVcache transfer: PUT, GET, and PUT_ASYNC. These methods can be specified using the `--kv-transfer-config` and `kv_connector_extra_config` parameters, specifically through the `send_type` field. Both PUT and PUT_ASYNC involve the P instance actively sending KVcache to the D instance. The difference is that PUT is a synchronous transfer method that blocks the main process, while PUT_ASYNC is an asynchronous transfer method. PUT_ASYNC uses a dedicated thread for sending KVcache, which means it does not block the main process. In contrast, the GET method involves the P instance saving the KVcache to the memory buffer after computing the prefill. The D instance then actively retrieves the computed KVcache from the P instance once it has allocated space for the KVcache.
34
+There are three methods for KVCache transfer: PUT, GET, and PUT_ASYNC. These methods can be specified using the `--kv-transfer-config` and `kv_connector_extra_config` parameters, specifically through the `send_type` field. Both PUT and PUT_ASYNC involve the P instance actively sending KVCache to the D instance. The difference is that PUT is a synchronous transfer method that blocks the main process, while PUT_ASYNC is an asynchronous transfer method. PUT_ASYNC uses a dedicated thread for sending KVCache, which means it does not block the main process. In contrast, the GET method involves the P instance saving the KVCache to the memory buffer after computing the prefill. The D instance then actively retrieves the computed KVCache from the P instance once it has allocated space for the KVCache.
35
36
Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
37
@@ -39,13 +39,13 @@ Experimental results have shown that the performance of these methods, from high
39
40
As long as the address of the counterpart is known, point-to-point KV cache transfer (using NCCL) can be performed, without being constrained by rank and world size. To support dynamic scaling (expansion and contraction) of instances with PD disaggregation. This means that adding or removing P/D instances does not require a full system restart.
41
42
-Each P/D instance only needs to create a single `P2pNcclEngine` instance. This instance maintains a ZMQ Server, which runs a dedicated thread to listen on the `zmq_addr` address and receive control flow requests from other instances. These requests include requests to establish an NCCL connection and requests to send KVcache metadata (such as tensor shapes and data types). However, it does not actually transmit the KVcache data itself.
42
+Each P/D instance only needs to create a single `P2pNcclEngine` instance. This instance maintains a ZMQ Server, which runs a dedicated thread to listen on the `zmq_addr` address and receive control flow requests from other instances. These requests include requests to establish an NCCL connection and requests to send KVCache metadata (such as tensor shapes and data types). However, it does not actually transmit the KVCache data itself.
43
44
-When a P instance and a D instance transmit KVcache for the first time, they need to establish a ZMQ connection and an NCCL group. For subsequent KVcache transmissions, this ZMQ connection and NCCL group are reused. The NCCL group consists of only two ranks, meaning the world size is equal to 2. This design is intended to support dynamic scaling, which means that adding or removing P/D instances does not require a full system restart. As long as the address of the counterpart is known, point-to-point KVcache transmission can be performed, without being restricted by rank or world size.
44
+When a P instance and a D instance transmit KVCache for the first time, they need to establish a ZMQ connection and an NCCL group. For subsequent KVCache transmissions, this ZMQ connection and NCCL group are reused. The NCCL group consists of only two ranks, meaning the world size is equal to 2. This design is intended to support dynamic scaling, which means that adding or removing P/D instances does not require a full system restart. As long as the address of the counterpart is known, point-to-point KVCache transmission can be performed, without being restricted by rank or world size.
45
46
## NCCL Group Topology
47
48
-Currently, only symmetric TP (Tensor Parallelism) methods are supported for KVcache transmission. Asymmetric TP and PP (Pipeline Parallelism) methods will be supported in the future. Figure 2 illustrates the 1P2D setup, where each instance has a TP (Tensor Parallelism) degree of 2. There are a total of 7 NCCL groups: three vLLM instances each have one NCCL group with TP=2. Additionally, the 0th GPU card of the P instance establishes an NCCL group with the 0th GPU card of each D instance. Similarly, the 1st GPU card of the P instance establishes an NCCL group with the 1st GPU card of each D instance.
48
+Currently, only symmetric TP (Tensor Parallelism) methods are supported for KVCache transmission. Asymmetric TP and PP (Pipeline Parallelism) methods will be supported in the future. Figure 2 illustrates the 1P2D setup, where each instance has a TP (Tensor Parallelism) degree of 2. There are a total of 7 NCCL groups: three vLLM instances each have one NCCL group with TP=2. Additionally, the 0th GPU card of the P instance establishes an NCCL group with the 0th GPU card of each D instance. Similarly, the 1st GPU card of the P instance establishes an NCCL group with the 1st GPU card of each D instance.
49
50
51
@@ -53,32 +53,18 @@ Each NCCL group occupies a certain amount of GPU memory buffer for communication
53
54
## GPU Memory Buffer and Tensor Memory Pool
55
56
-The trade-off in the size of the memory buffer is as follows: For P instances, the memory buffer is not required in PUT and PUT_ASYNC modes, but it is necessary in GET mode. For D instances, a memory buffer is needed in all three modes. The memory buffer for D instances should not be too large. Similarly, for P instances in GET mode, the memory buffer should also not be too large. The memory buffer of D instances is used to temporarily store KVcache sent by P instances. If it is too large, it will reduce the KVcache space available for normal inference by D instances, thereby decreasing the inference batch size and ultimately leading to a reduction in output throughput. The size of the memory buffer is configured by the parameter `kv_buffer_size`, measured in bytes, and is typically set to 5%～10% of the memory size.
56
+The trade-off in the size of the memory buffer is as follows: For P instances, the memory buffer is not required in PUT and PUT_ASYNC modes, but it is necessary in GET mode. For D instances, a memory buffer is needed in all three modes. The memory buffer for D instances should not be too large. Similarly, for P instances in GET mode, the memory buffer should also not be too large. The memory buffer of D instances is used to temporarily store KVCache sent by P instances. If it is too large, it will reduce the KVCache space available for normal inference by D instances, thereby decreasing the inference batch size and ultimately leading to a reduction in output throughput. The size of the memory buffer is configured by the parameter `kv_buffer_size`, measured in bytes, and is typically set to 5%～10% of the memory size.
57
58
-If the `--max-num-seqs` parameter for P instances is set to a large value, due to the large batch size, P instances will generate a large amount of KVcache simultaneously. This may exceed the capacity of the memory buffer of D instances, resulting in KVcache loss. Once KVcache is lost, D instances need to recompute Prefill, which is equivalent to performing Prefill twice. Consequently, the time-to-first-token (TTFT) will significantly increase, leading to degraded performance.
58
+If the `--max-num-seqs` parameter for P instances is set to a large value, due to the large batch size, P instances will generate a large amount of KVCache simultaneously. This may exceed the capacity of the memory buffer of D instances, resulting in KVCache loss. Once KVCache is lost, D instances need to recompute Prefill, which is equivalent to performing Prefill twice. Consequently, the time-to-first-token (TTFT) will significantly increase, leading to degraded performance.
59
60
-To address the above issues, I have designed and developed a local Tensor memory pool for storing KVcache, inspired by the buddy system used in Linux memory modules. Since the memory is sufficiently large, typically in the TB range on servers, there is no need to consider prefix caching or using block-based designs to reuse memory, thereby saving space. When the memory buffer is insufficient, KVcache can be directly stored in the Tensor memory pool, and D instances can subsequently retrieve KVcache from it. The read and write speed is that of PCIe, with PCIe 4.0 having a speed of approximately 21 GB/s, which is usually faster than the Prefill speed. Otherwise, solutions like Mooncake and lmcache would not be necessary. The Tensor memory pool acts as a flood diversion area, typically unused except during sudden traffic surges. In the worst-case scenario, my solution performs no worse than the normal situation with a Cache store.
60
+To address the above issues, I have designed and developed a local Tensor memory pool for storing KVCache, inspired by the buddy system used in Linux memory modules. Since the memory is sufficiently large, typically in the TB range on servers, there is no need to consider prefix caching or using block-based designs to reuse memory, thereby saving space. When the memory buffer is insufficient, KVCache can be directly stored in the Tensor memory pool, and D instances can subsequently retrieve KVCache from it. The read and write speed is that of PCIe, with PCIe 4.0 having a speed of approximately 21 GB/s, which is usually faster than the Prefill speed. Otherwise, solutions like Mooncake and lmcache would not be necessary. The Tensor memory pool acts as a flood diversion area, typically unused except during sudden traffic surges. In the worst-case scenario, my solution performs no worse than the normal situation with a Cache store.
61
62
# Install vLLM
63
64
??? console "Commands"
65
66
```shell
67
- # Enter the home directory or your working directory.
68
- cd /home
69
-
70
- # Download the installation package, and I will update the commit-id in time. You can directly copy the command.
71
- wget https://vllm-wheels.s3.us-west-2.amazonaws.com/9112b443a042d8d815880b8780633882ad32b183/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
72
-
73
- # Download the code repository.
74
- git clone -b xpyd-v1 https://github.com/Abatom/vllm.git
75
- cd vllm
76
-
77
- # Set the installation package path.
78
- export VLLM_PRECOMPILED_WHEEL_LOCATION=/home/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
79
-
80
- # installation
81
- pip install -e . -v
67
+ pip install "vllm>=0.9.2"
68
```
69
70
# Run xPyD
@@ -90,7 +76,7 @@ To address the above issues, I have designed and developed a local Tensor memory
76
- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
77
- `PUT_ASYNC` offers the best performance and should be prioritized.
78
- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
93
-- The `disagg_prefill_proxy_xpyd.py` script will use port 10001 (for receiving client requests) and port 30001 (for receiving service discovery from P and D instances).
79
+- The `disagg_proxy_p2p_nccl_xpyd.py` script will use port 10001 (for receiving client requests) and port 30001 (for receiving service discovery from P and D instances).
80
- The node running the proxy must have `quart` installed.
81
- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
82
- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
@@ -100,8 +86,8 @@ To address the above issues, I have designed and developed a local Tensor memory
86
### Proxy (e.g. 10.0.1.1)
87
88
```shell
103
-cd {your vllm directory}/examples/online_serving/disagg_xpyd/
104
-python3 disagg_prefill_proxy_xpyd.py &
89
+cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
90
+python3 disagg_proxy_p2p_nccl_xpyd.py &
91
```
92
93
### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
@@ -111,7 +97,7 @@ python3 disagg_prefill_proxy_xpyd.py &
97
```shell
98
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
99
--host 0.0.0.0 \
114
- --port 20005 \
100
+ --port 20001 \
101
--tensor-parallel-size 1 \
102
--seed 1024 \
103
--served-model-name base_model \
@@ -123,7 +109,7 @@ python3 disagg_prefill_proxy_xpyd.py &
109
--gpu-memory-utilization 0.9 \
110
--disable-log-request \
111
--kv-transfer-config \
126
- '{"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":"20005","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
112
+ '{"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 &
113
```
114
115
### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
@@ -133,7 +119,7 @@ python3 disagg_prefill_proxy_xpyd.py &
119
```shell
120
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
121
--host 0.0.0.0 \
136
- --port 20009 \
122
+ --port 20002 \
123
--tensor-parallel-size 1 \
124
--seed 1024 \
125
--served-model-name base_model \
@@ -145,7 +131,7 @@ python3 disagg_prefill_proxy_xpyd.py &
131
--gpu-memory-utilization 0.7 \
132
--disable-log-request \
133
--kv-transfer-config \
148
- '{"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":"20009","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
134
+ '{"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 &
135
```
136
137
### Decode2 (e.g. 10.0.1.4 or 10.0.1.1)
@@ -167,7 +153,7 @@ python3 disagg_prefill_proxy_xpyd.py &
153
--gpu-memory-utilization 0.7 \
154
--disable-log-request \
155
--kv-transfer-config \
170
- '{"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","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
156
+ '{"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 &
157
```
158
159
### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
@@ -177,7 +163,7 @@ python3 disagg_prefill_proxy_xpyd.py &
163
```shell
164
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
165
--host 0.0.0.0 \
180
- --port 20008 \
166
+ --port 20004 \
167
--tensor-parallel-size 1 \
168
--seed 1024 \
169
--served-model-name base_model \
@@ -189,7 +175,7 @@ python3 disagg_prefill_proxy_xpyd.py &
175
--gpu-memory-utilization 0.7 \
176
--disable-log-request \
177
--kv-transfer-config \
192
- '{"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":"20008","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
178
+ '{"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 &
179
```
180
181
## Run 3P1D
@@ -197,8 +183,8 @@ python3 disagg_prefill_proxy_xpyd.py &
183
### Proxy (e.g. 10.0.1.1)
184
185
```shell
200
-cd {your vllm directory}/examples/online_serving/disagg_xpyd/
201
-python3 disagg_prefill_proxy_xpyd.py &
186
+cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
187
+python3 disagg_proxy_p2p_nccl_xpyd.py &
188
```
189
190
### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
@@ -208,7 +194,7 @@ python3 disagg_prefill_proxy_xpyd.py &
194
```shell
195
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
196
--host 0.0.0.0 \
211
- --port 20005 \
197
+ --port 20001 \
198
--tensor-parallel-size 1 \
199
--seed 1024 \
200
--served-model-name base_model \
@@ -220,7 +206,7 @@ python3 disagg_prefill_proxy_xpyd.py &
206
--gpu-memory-utilization 0.9 \
207
--disable-log-request \
208
--kv-transfer-config \
223
- '{"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":"20005","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
209
+ '{"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 &
210
```
211
212
### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
@@ -230,7 +216,7 @@ python3 disagg_prefill_proxy_xpyd.py &
216
```shell
217
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
218
--host 0.0.0.0 \
233
- --port 20009 \
219
+ --port 20002 \
220
--tensor-parallel-size 1 \
221
--seed 1024 \
222
--served-model-name base_model \
@@ -242,7 +228,7 @@ python3 disagg_prefill_proxy_xpyd.py &
228
--gpu-memory-utilization 0.9 \
229
--disable-log-request \
230
--kv-transfer-config \
245
- '{"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":"20009","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
231
+ '{"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 &
232
```
233
234
### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
@@ -264,7 +250,7 @@ python3 disagg_prefill_proxy_xpyd.py &
250
--gpu-memory-utilization 0.9 \
251
--disable-log-request \
252
--kv-transfer-config \
267
- '{"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","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
253
+ '{"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 &
254
```
255
256
### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
@@ -274,7 +260,7 @@ python3 disagg_prefill_proxy_xpyd.py &
260
```shell
261
VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
262
--host 0.0.0.0 \
277
- --port 20008 \
263
+ --port 20004 \
264
--tensor-parallel-size 1 \
265
--seed 1024 \
266
--served-model-name base_model \
@@ -286,7 +272,7 @@ python3 disagg_prefill_proxy_xpyd.py &
272
--gpu-memory-utilization 0.7 \
273
--disable-log-request \
274
--kv-transfer-config \
289
- '{"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":"20008","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
275
+ '{"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 &
276
```
277
278
# Single request
@@ -334,24 +320,6 @@ pgrep python | xargs kill -9 && pkill -f python
320
321
# Test data
322
337
-## **Scenario 1**: 1K input & 1K output tokens, E2E P99 latency ~20s
338
-- **1P5D (6×A800) vs vLLM (1×A800)**:
339
- - Throughput ↑7.2% (1085 → 6979/6)
340
- - ITL (P99) ↓81.3% (120ms → 22.9ms)
341
- - TTFT (P99) ↑26.8% (175ms → 222ms)
342
- - TPOT: No change
343
-
344
-- **1P6D (7×A800) vs vLLM (1×A800)**:
345
- - Throughput ↑9.6% (1085 → 8329/7)
346
- - ITL (P99) ↓81.0% (120ms → 22.7ms)
347
- - TTFT (P99) ↑210% (175ms →543ms)
348
- - TPOT: No change
349
-
350
-## **Scenario 2**: 1K input & 200 output tokens, E2E P99 latency ~4s
351
-- **1P1D (2×A800) vs vLLM (1×A800)**:
352
- - Throughput ↑37.4% (537 → 1476/2)
353
- - ITL (P99) ↓81.8% (127ms → 23.1ms)
354
- - TTFT (P99) ↑41.8% (160ms → 227ms)
355
- - TPOT: No change
356
-
357
-
323
+## **Scenario**: 1K input & 200 output tokens, E2E P99 latency ~2s
324
+
325
+
examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_proxy_p2p_nccl_xpyd.py
+35
-4
@@ -4,7 +4,9 @@
4
import os
5
import socket
6
import threading
7
+import time
8
import uuid
9
+from typing import Any
10
11
import aiohttp
12
import msgpack
@@ -12,12 +14,25 @@ import zmq
14
from quart import Quart, make_response, request
15
16
count = 0
15
-prefill_instances: dict[str, str] = {} # http_address: zmq_address
16
-decode_instances: dict[str, str] = {} # http_address: zmq_address
17
+prefill_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
18
+decode_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
19
20
prefill_cv = threading.Condition()
21
decode_cv = threading.Condition()
22
23
+DEFAULT_PING_SECONDS = 5
24
+
25
+
26
+def _remove_oldest_instances(instances: dict[str, Any]) -> None:
27
+ oldest_key = next(iter(instances), None)
28
+ while oldest_key is not None:
29
+ value = instances[oldest_key]
30
+ if value[1] > time.time():
31
+ break
32
+ print(f"🔴Remove [HTTP:{oldest_key}, ZMQ:{value[0]}, stamp:{value[1]}]")
33
+ instances.pop(oldest_key, None)
34
+ oldest_key = next(iter(instances), None)
35
+
36
37
def _listen_for_register(poller, router_socket):
38
while True:
@@ -31,12 +46,23 @@ def _listen_for_register(poller, router_socket):
46
global prefill_instances
47
global prefill_cv
48
with prefill_cv:
34
- prefill_instances[data["http_address"]] = data["zmq_address"]
49
+ node = prefill_instances.pop(data["http_address"], None)
50
+ prefill_instances[data["http_address"]] = (
51
+ data["zmq_address"],
52
+ time.time() + DEFAULT_PING_SECONDS,
53
+ )
54
+ _remove_oldest_instances(prefill_instances)
55
+
56
elif data["type"] == "D":
57
global decode_instances
58
global decode_cv
59
with decode_cv:
39
- decode_instances[data["http_address"]] = data["zmq_address"]
60
+ node = decode_instances.pop(data["http_address"], None)
61
+ decode_instances[data["http_address"]] = (
62
+ data["zmq_address"],
63
+ time.time() + DEFAULT_PING_SECONDS,
64
+ )
65
+ _remove_oldest_instances(decode_instances)
66
else:
67
print(
68
"Unexpected, Received message from %s, data: %s",
@@ -44,6 +70,9 @@ def _listen_for_register(poller, router_socket):
70
data,
71
)
72
73
+ if node is None:
74
+ print(f"🔵Add [HTTP:{data['http_address']}, ZMQ:{data['zmq_address']}]")
75
+
76
77
def start_service_discovery(hostname, port):
78
if not hostname:
@@ -105,12 +134,14 @@ async def handle_request():
134
with prefill_cv:
135
prefill_list = list(prefill_instances.items())
136
prefill_addr, prefill_zmq_addr = prefill_list[count % len(prefill_list)]
137
+ prefill_zmq_addr = prefill_zmq_addr[0]
138
139
global decode_instances
140
global decode_cv
141
with decode_cv:
142
decode_list = list(decode_instances.items())
143
decode_addr, decode_zmq_addr = decode_list[count % len(decode_list)]
144
+ decode_zmq_addr = decode_zmq_addr[0]
145
146
print(
147
f"handle_request count: {count}, [HTTP:{prefill_addr}, "
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+7
-31
@@ -13,7 +13,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
13
from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
14
P2pNcclEngine)
15
from vllm.distributed.parallel_state import get_world_group
16
-from vllm.forward_context import get_forward_context
16
from vllm.logger import init_logger
17
from vllm.v1.attention.backends.mla.common import MLACommonMetadata
18
from vllm.v1.core.sched.output import SchedulerOutput
@@ -238,32 +237,16 @@ class P2pNcclConnector(KVConnectorBase_V1):
237
238
assert self.p2p_nccl_engine is not None
239
241
- def extract_kv_from_layer(
242
- layer: torch.Tensor,
243
- slot_mapping: torch.Tensor,
244
- ) -> torch.Tensor:
245
- """Extract the KV cache from the layer.
246
-
247
- Assume the shape of the layer is (2, num_pages, page_size, xxx)
248
- if MLA is not used, and (num_pages, page_size, xxx) otherwise.
249
- """
250
- if isinstance(attn_metadata, MLACommonMetadata):
251
- num_pages, page_size = layer.shape[0], layer.shape[1]
252
- return layer.reshape(num_pages * page_size, -1)[slot_mapping,
253
- ...]
254
- num_pages, page_size = layer.shape[1], layer.shape[2]
255
- return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
256
- ...]
257
-
240
connector_metadata = self._get_connector_metadata()
241
assert isinstance(connector_metadata, P2pNcclConnectorMetadata)
242
for request in connector_metadata.requests:
243
request_id = request.request_id
244
ip, port = self.parse_request_id(request_id, True)
245
remote_address = ip + ":" + str(port + self._rank)
264
- kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
265
- self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
266
- kv_cache, remote_address)
246
+ self.p2p_nccl_engine.send_tensor(
247
+ request_id + "#" + layer_name, kv_layer, remote_address,
248
+ request.slot_mapping,
249
+ isinstance(attn_metadata, MLACommonMetadata))
250
251
def wait_for_save(self):
252
if self.is_producer:
@@ -286,9 +269,10 @@ class P2pNcclConnector(KVConnectorBase_V1):
269
270
assert self.p2p_nccl_engine is not None
271
289
- forward_context: ForwardContext = get_forward_context()
272
+ no_compile_layers = (
273
+ self._vllm_config.compilation_config.static_forward_context)
274
return self.p2p_nccl_engine.get_finished(finished_req_ids,
291
- forward_context)
275
+ no_compile_layers)
276
277
# ==============================
278
# Scheduler-side methods
@@ -418,14 +402,6 @@ class P2pNcclConnector(KVConnectorBase_V1):
402
block_ids=block_ids,
403
block_size=self._block_size)
404
421
- # Requests loaded asynchronously are not in the scheduler_output.
422
- # for request_id in self._requests_need_load:
423
- # request, block_ids = self._requests_need_load[request_id]
424
- # meta.add_request(request_id=request.request_id,
425
- # token_ids=request.prompt_token_ids,
426
- # block_ids=block_ids,
427
- # block_size=self._block_size)
428
-
405
self._requests_need_load.clear()
406
return meta
407
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+194
-159
@@ -8,7 +8,8 @@ import time
8
import typing
9
from collections import deque
10
from contextlib import contextmanager
11
-from typing import TYPE_CHECKING, Any, Optional
11
+from dataclasses import dataclass
12
+from typing import Any, Optional
13
14
import msgpack
15
import torch
@@ -21,9 +22,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import
22
TensorMemoryPool)
23
from vllm.utils import current_stream, get_ip
24
24
-if TYPE_CHECKING:
25
- from vllm.forward_context import ForwardContext
26
-
25
logger = logging.getLogger(__name__)
26
27
DEFAULT_MEM_POOL_SIZE_GB = 32
@@ -59,6 +57,15 @@ def set_p2p_nccl_context(num_channels: str):
57
os.environ.pop(var, None)
58
59
60
+@dataclass
61
+class SendQueueItem:
62
+ tensor_id: str
63
+ remote_address: str
64
+ tensor: torch.Tensor
65
+ slot_mapping: torch.Tensor
66
+ is_mla: bool
67
+
68
+
69
class P2pNcclEngine:
70
71
def __init__(self,
@@ -112,24 +119,26 @@ class P2pNcclEngine:
119
self.send_stream = torch.cuda.Stream()
120
self.recv_stream = torch.cuda.Stream()
121
115
- mem_pool_size_gb = self.config.get_from_extra_config(
116
- "mem_pool_size_gb", DEFAULT_MEM_POOL_SIZE_GB)
117
- self.pool = TensorMemoryPool(max_block_size=int(mem_pool_size_gb) *
118
- 1024**3) # GB
122
+ mem_pool_size_gb = float(
123
+ self.config.get_from_extra_config("mem_pool_size_gb",
124
+ DEFAULT_MEM_POOL_SIZE_GB))
125
+ self.pool = TensorMemoryPool(max_block_size=int(mem_pool_size_gb *
126
+ 1024**3)) # GB
127
128
# The sending type includes tree mutually exclusive options:
129
# PUT, GET, PUT_ASYNC.
122
- self.send_type = self.config.get_from_extra_config("send_type", "PUT")
130
+ self.send_type = self.config.get_from_extra_config(
131
+ "send_type", "PUT_ASYNC")
132
if self.send_type == "GET":
133
# tensor_id: torch.Tensor
134
self.send_store: dict[str, torch.Tensor] = {}
135
else:
136
# PUT or PUT_ASYNC
137
# tensor_id: torch.Tensor
129
- self.send_queue: deque[list[Any]] = deque()
138
+ self.send_queue: deque[SendQueueItem] = deque()
139
self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
140
if self.send_type == "PUT_ASYNC":
132
- self._send_thread = threading.Thread(target=self._send_async,
141
+ self._send_thread = threading.Thread(target=self.send_async,
142
daemon=True)
143
self._send_thread.start()
144
@@ -146,13 +155,12 @@ class P2pNcclEngine:
155
"nccl_num_channels", "8")
156
157
self._listener_thread = threading.Thread(
149
- target=self._listen_for_requests, daemon=True)
158
+ target=self.listen_for_requests, daemon=True)
159
self._listener_thread.start()
160
161
self._ping_thread = None
162
if port_offset == 0 and self.proxy_address != "":
154
- self._ping_thread = threading.Thread(target=self._ping,
155
- daemon=True)
163
+ self._ping_thread = threading.Thread(target=self.ping, daemon=True)
164
self._ping_thread.start()
165
166
logger.info(
@@ -162,7 +170,7 @@ class P2pNcclEngine:
170
self.http_address, self.zmq_address, self.proxy_address,
171
self.send_type, self.buffer_size_threshold, self.nccl_num_channels)
172
165
- def _create_connect(self, remote_address: typing.Optional[str] = None):
173
+ def create_connect(self, remote_address: typing.Optional[str] = None):
174
assert remote_address is not None
175
if remote_address not in self.socks:
176
sock = self.context.socket(zmq.DEALER)
@@ -184,7 +192,7 @@ class P2pNcclEngine:
192
comm: ncclComm_t = self.nccl.ncclCommInitRank(
193
2, unique_id, rank)
194
self.comms[remote_address] = (comm, rank)
187
- logger.info("🤝ncclCommInitRank Success, %s👉%s, MyRank: %s",
195
+ logger.info("🤝ncclCommInitRank Success, %s👉%s, MyRank:%s",
196
self.zmq_address, remote_address, rank)
197
198
return self.socks[remote_address], self.comms[remote_address]
@@ -194,44 +202,54 @@ class P2pNcclEngine:
202
tensor_id: str,
203
tensor: torch.Tensor,
204
remote_address: typing.Optional[str] = None,
205
+ slot_mapping: torch.Tensor = None,
206
+ is_mla: bool = False,
207
) -> bool:
208
if remote_address is None:
209
with self.recv_store_cv:
210
self.recv_store[tensor_id] = tensor
211
self.recv_store_cv.notify()
212
return True
203
- else:
204
- if self.send_type == "PUT":
205
- return self._send_sync(tensor_id, tensor, remote_address)
206
- elif self.send_type == "PUT_ASYNC":
207
- with self.send_queue_cv:
208
- self.send_queue.append([tensor_id, remote_address, tensor])
209
- self.send_queue_cv.notify()
210
- else: # GET
211
- with self.send_store_cv:
212
- tensor_size = tensor.element_size() * tensor.numel()
213
- while (self.buffer_size + tensor_size
214
- > self.buffer_size_threshold):
215
- oldest_tenser_id = next(iter(self.send_store))
216
- oldest_tenser = self.send_store.pop(oldest_tenser_id)
217
- oldest_tenser_size = oldest_tenser.element_size(
218
- ) * oldest_tenser.numel()
219
- self.buffer_size -= oldest_tenser_size
220
- logger.info(
221
- "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
222
- " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
223
- remote_address, tensor_id, tensor_size,
224
- self.buffer_size, oldest_tenser_size, self.rank)
225
-
226
- self.send_store[tensor_id] = tensor
227
- self.buffer_size += tensor_size
228
- logger.debug(
229
- "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
230
- "shape:%s, rank:%d, buffer_size:%d(%.2f%%)",
231
- remote_address, tensor_id, tensor_size, tensor.shape,
232
- self.rank, self.buffer_size,
233
- self.buffer_size / self.buffer_size_threshold * 100)
213
214
+ item = SendQueueItem(tensor_id=tensor_id,
215
+ remote_address=remote_address,
216
+ tensor=tensor,
217
+ slot_mapping=slot_mapping,
218
+ is_mla=is_mla)
219
+
220
+ if self.send_type == "PUT":
221
+ return self.send_sync(item)
222
+
223
+ if self.send_type == "PUT_ASYNC":
224
+ with self.send_queue_cv:
225
+ self.send_queue.append(item)
226
+ self.send_queue_cv.notify()
227
+ return True
228
+
229
+ # GET
230
+ with self.send_store_cv:
231
+ tensor_size = tensor.element_size() * tensor.numel()
232
+ while (self.buffer_size + tensor_size
233
+ > self.buffer_size_threshold):
234
+ oldest_tenser_id = next(iter(self.send_store))
235
+ oldest_tenser = self.send_store.pop(oldest_tenser_id)
236
+ oldest_tenser_size = oldest_tenser.element_size(
237
+ ) * oldest_tenser.numel()
238
+ self.buffer_size -= oldest_tenser_size
239
+ logger.info(
240
+ "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
241
+ " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
242
+ remote_address, tensor_id, tensor_size, self.buffer_size,
243
+ oldest_tenser_size, self.rank)
244
+
245
+ self.send_store[tensor_id] = tensor
246
+ self.buffer_size += tensor_size
247
+ logger.debug(
248
+ "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
249
+ "shape:%s, rank:%d, buffer_size:%d(%.2f%%)", remote_address,
250
+ tensor_id, tensor_size, tensor.shape, self.rank,
251
+ self.buffer_size,
252
+ self.buffer_size / self.buffer_size_threshold * 100)
253
return True
254
255
def recv_tensor(
@@ -267,7 +285,7 @@ class P2pNcclEngine:
285
return None
286
287
if remote_address not in self.socks:
270
- self._create_connect(remote_address)
288
+ self.create_connect(remote_address)
289
290
sock = self.socks[remote_address]
291
comm, rank = self.comms[remote_address]
@@ -282,121 +300,121 @@ class P2pNcclEngine:
300
remote_address, tensor_id, data["ret"])
301
return None
302
285
- tensor = torch.empty(data["shape"],
286
- dtype=getattr(torch, data["dtype"]),
287
- device=self.device)
303
+ with torch.cuda.stream(self.recv_stream):
304
+ tensor = torch.empty(data["shape"],
305
+ dtype=getattr(torch, data["dtype"]),
306
+ device=self.device)
307
289
- self._recv(comm, tensor, rank ^ 1, self.recv_stream)
308
+ self.recv(comm, tensor, rank ^ 1, self.recv_stream)
309
310
return tensor
311
293
- def _listen_for_requests(self):
312
+ def listen_for_requests(self):
313
while True:
314
socks = dict(self.poller.poll())
296
- if self.router_socket in socks:
297
- remote_address, message = self.router_socket.recv_multipart()
298
- data = msgpack.loads(message)
299
- if data["cmd"] == "NEW":
300
- unique_id = self.nccl.unique_id_from_bytes(
301
- bytes(data["unique_id"]))
302
- with torch.cuda.device(self.device):
303
- rank = 1
304
- with set_p2p_nccl_context(self.nccl_num_channels):
305
- comm: ncclComm_t = self.nccl.ncclCommInitRank(
306
- 2, unique_id, rank)
307
- self.comms[remote_address.decode()] = (comm, rank)
308
- logger.info(
309
- "🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
310
- self.zmq_address, remote_address.decode(), rank)
311
- elif data["cmd"] == "PUT":
312
- tensor_id = data["tensor_id"]
313
- try:
314
- with torch.cuda.stream(self.recv_stream):
315
- tensor = torch.empty(data["shape"],
316
- dtype=getattr(
317
- torch, data["dtype"]),
318
- device=self.device)
319
- self.router_socket.send_multipart(
320
- [remote_address, b"0"])
321
- comm, rank = self.comms[remote_address.decode()]
322
- self._recv(comm, tensor, rank ^ 1, self.recv_stream)
323
- tensor_size = tensor.element_size() * tensor.numel()
324
- if (self.buffer_size + tensor_size
325
- > self.buffer_size_threshold):
326
- # Store Tensor in memory pool
327
- addr = self.pool.store_tensor(tensor)
328
- tensor = (addr, tensor.dtype, tensor.shape)
329
- logger.warning(
330
- "🔴[PUT]Recv Tensor, Out Of Threshold, "
331
- "%s👈%s, data:%s, addr:%d", self.zmq_address,
332
- remote_address.decode(), data, addr)
333
- else:
334
- self.buffer_size += tensor_size
335
-
336
- except torch.cuda.OutOfMemoryError:
337
- self.router_socket.send_multipart(
338
- [remote_address, b"1"])
339
- tensor = None
315
+ if self.router_socket not in socks:
316
+ continue
317
+
318
+ remote_address, message = self.router_socket.recv_multipart()
319
+ data = msgpack.loads(message)
320
+ if data["cmd"] == "NEW":
321
+ unique_id = self.nccl.unique_id_from_bytes(
322
+ bytes(data["unique_id"]))
323
+ with torch.cuda.device(self.device):
324
+ rank = 1
325
+ with set_p2p_nccl_context(self.nccl_num_channels):
326
+ comm: ncclComm_t = self.nccl.ncclCommInitRank(
327
+ 2, unique_id, rank)
328
+ self.comms[remote_address.decode()] = (comm, rank)
329
+ logger.info("🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
330
+ self.zmq_address, remote_address.decode(),
331
+ rank)
332
+ elif data["cmd"] == "PUT":
333
+ tensor_id = data["tensor_id"]
334
+ try:
335
+ with torch.cuda.stream(self.recv_stream):
336
+ tensor = torch.empty(data["shape"],
337
+ dtype=getattr(
338
+ torch, data["dtype"]),
339
+ device=self.device)
340
+ self.router_socket.send_multipart([remote_address, b"0"])
341
+ comm, rank = self.comms[remote_address.decode()]
342
+ self.recv(comm, tensor, rank ^ 1, self.recv_stream)
343
+ tensor_size = tensor.element_size() * tensor.numel()
344
+ if (self.buffer_size + tensor_size
345
+ > self.buffer_size_threshold):
346
+ # Store Tensor in memory pool
347
+ addr = self.pool.store_tensor(tensor)
348
+ tensor = (addr, tensor.dtype, tensor.shape)
349
logger.warning(
341
- "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, "
342
- "data:%s", self.zmq_address,
343
- remote_address.decode(), data)
344
-
345
- with self.recv_store_cv:
346
- self.recv_store[tensor_id] = tensor
347
- self._have_received_tensor_id(tensor_id)
348
- self.recv_store_cv.notify()
349
-
350
- elif data["cmd"] == "GET":
351
- tensor_id = data["tensor_id"]
352
- with self.send_store_cv:
353
- tensor = self.send_store.pop(tensor_id, None)
354
- if tensor is not None:
355
- data = {
356
- "ret": 0,
357
- "shape": tensor.shape,
358
- "dtype":
359
- str(tensor.dtype).replace("torch.", "")
360
- }
361
- # LRU
362
- self.send_store[tensor_id] = tensor
363
- self._have_sent_tensor_id(tensor_id)
364
- else:
365
- data = {"ret": 1}
366
-
367
- self.router_socket.send_multipart(
368
- [remote_address, msgpack.dumps(data)])
369
-
370
- if data["ret"] == 0:
371
- comm, rank = self.comms[remote_address.decode()]
372
- self._send(comm, tensor.to(self.device), rank ^ 1,
373
- self.send_stream)
374
- else:
350
+ "🔴[PUT]Recv Tensor, Out Of Threshold, "
351
+ "%s👈%s, data:%s, addr:%d", self.zmq_address,
352
+ remote_address.decode(), data, addr)
353
+ else:
354
+ self.buffer_size += tensor_size
355
+
356
+ except torch.cuda.OutOfMemoryError:
357
+ self.router_socket.send_multipart([remote_address, b"1"])
358
+ tensor = None
359
logger.warning(
376
- "🚧Unexpected, Received message from %s, data:%s",
377
- remote_address, data)
360
+ "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, "
361
+ "data:%s", self.zmq_address, remote_address.decode(),
362
+ data)
363
379
- def _have_sent_tensor_id(self, tensor_id: str):
364
+ with self.recv_store_cv:
365
+ self.recv_store[tensor_id] = tensor
366
+ self.have_received_tensor_id(tensor_id)
367
+ self.recv_store_cv.notify()
368
+
369
+ elif data["cmd"] == "GET":
370
+ tensor_id = data["tensor_id"]
371
+ with self.send_store_cv:
372
+ tensor = self.send_store.pop(tensor_id, None)
373
+ if tensor is not None:
374
+ data = {
375
+ "ret": 0,
376
+ "shape": tensor.shape,
377
+ "dtype": str(tensor.dtype).replace("torch.", "")
378
+ }
379
+ # LRU
380
+ self.send_store[tensor_id] = tensor
381
+ self.have_sent_tensor_id(tensor_id)
382
+ else:
383
+ data = {"ret": 1}
384
+
385
+ self.router_socket.send_multipart(
386
+ [remote_address, msgpack.dumps(data)])
387
+
388
+ if data["ret"] == 0:
389
+ comm, rank = self.comms[remote_address.decode()]
390
+ self.send(comm, tensor.to(self.device), rank ^ 1,
391
+ self.send_stream)
392
+ else:
393
+ logger.warning(
394
+ "🚧Unexpected, Received message from %s, data:%s",
395
+ remote_address, data)
396
+
397
+ def have_sent_tensor_id(self, tensor_id: str):
398
request_id = tensor_id.split('#')[0]
399
if request_id not in self.send_request_id_to_tensor_ids:
400
self.send_request_id_to_tensor_ids[request_id] = set()
401
self.send_request_id_to_tensor_ids[request_id].add(tensor_id)
402
385
- def _have_received_tensor_id(self, tensor_id: str):
403
+ def have_received_tensor_id(self, tensor_id: str):
404
request_id = tensor_id.split('#')[0]
405
if request_id not in self.recv_request_id_to_tensor_ids:
406
self.recv_request_id_to_tensor_ids[request_id] = set()
407
self.recv_request_id_to_tensor_ids[request_id].add(tensor_id)
408
391
- def _send_async(self):
409
+ def send_async(self):
410
while True:
411
with self.send_queue_cv:
412
while not self.send_queue:
413
self.send_queue_cv.wait()
396
- tensor_id, remote_address, tensor = self.send_queue.popleft()
414
+ item = self.send_queue.popleft()
415
if not self.send_queue:
416
self.send_queue_cv.notify()
399
- self._send_sync(tensor_id, tensor, remote_address)
417
+ self.send_sync(item)
418
419
def wait_for_sent(self):
420
if self.send_type == "PUT_ASYNC":
@@ -409,22 +427,21 @@ class P2pNcclEngine:
427
"🚧[PUT_ASYNC]It took %.3fms to wait for the send_queue"
428
" to be empty, rank:%d", duration * 1000, self.rank)
429
412
- def _send_sync(
413
- self,
414
- tensor_id: str,
415
- tensor: torch.Tensor,
416
- remote_address: typing.Optional[str] = None,
417
- ) -> bool:
418
- if remote_address is None:
430
+ def send_sync(self, item: SendQueueItem) -> bool:
431
+ if item.remote_address is None:
432
return False
420
- if remote_address not in self.socks:
421
- self._create_connect(remote_address)
433
+ if item.remote_address not in self.socks:
434
+ self.create_connect(item.remote_address)
435
423
- sock = self.socks[remote_address]
424
- comm, rank = self.comms[remote_address]
436
+ with self.send_stream:
437
+ tensor = self.extract_kv_from_layer(item.is_mla, item.tensor,
438
+ item.slot_mapping)
439
+
440
+ sock = self.socks[item.remote_address]
441
+ comm, rank = self.comms[item.remote_address]
442
data = {
443
"cmd": "PUT",
427
- "tensor_id": tensor_id,
444
+ "tensor_id": item.tensor_id,
445
"shape": tensor.shape,
446
"dtype": str(tensor.dtype).replace("torch.", "")
447
}
@@ -435,20 +452,21 @@ class P2pNcclEngine:
452
logger.error(
453
"🔴Send Tensor, Peer Out Of Memory/Threshold, %s 👉 %s, "
454
"MyRank:%s, data:%s, tensor:%s, size:%fGB, response:%s",
438
- self.zmq_address, remote_address, rank, data, tensor.shape,
455
+ self.zmq_address, item.remote_address, rank, data,
456
+ tensor.shape,
457
tensor.element_size() * tensor.numel() / 1024**3,
458
response.decode())
459
return False
460
443
- self._send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
461
+ self.send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
462
463
if self.send_type == "PUT_ASYNC":
446
- self._have_sent_tensor_id(tensor_id)
464
+ self.have_sent_tensor_id(item.tensor_id)
465
466
return True
467
468
def get_finished(
451
- self, finished_req_ids: set[str], forward_context: "ForwardContext"
469
+ self, finished_req_ids: set[str], no_compile_layers
470
) -> tuple[Optional[set[str]], Optional[set[str]]]:
471
"""
472
Notifies worker-side connector ids of requests that have
@@ -463,7 +481,7 @@ class P2pNcclEngine:
481
482
# Clear the buffer upon request completion.
483
for request_id in finished_req_ids:
466
- for layer_name in forward_context.no_compile_layers:
484
+ for layer_name in no_compile_layers:
485
tensor_id = request_id + "#" + layer_name
486
if tensor_id in self.recv_store:
487
with self.recv_store_cv:
@@ -472,7 +490,6 @@ class P2pNcclEngine:
490
request_id, None)
491
self.recv_request_id_to_tensor_ids.pop(
492
request_id, None)
475
- addr = 0
493
if isinstance(tensor, tuple):
494
addr, _, _ = tensor
495
self.pool.free(addr)
@@ -485,7 +502,7 @@ class P2pNcclEngine:
502
503
return finished_sending or None, finished_recving or None
504
488
- def _ping(self):
505
+ def ping(self):
506
sock = self.context.socket(zmq.DEALER)
507
sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
508
logger.debug("ping start, zmq_address:%s", self.zmq_address)
@@ -499,7 +516,7 @@ class P2pNcclEngine:
516
sock.send(msgpack.dumps(data))
517
time.sleep(3)
518
502
- def _send(self, comm, tensor: torch.Tensor, dst: int, stream=None):
519
+ def send(self, comm, tensor: torch.Tensor, dst: int, stream=None):
520
assert tensor.device == self.device, (
521
f"this nccl communicator is created to work on {self.device}, "
522
f"but the input tensor is on {tensor.device}")
@@ -512,7 +529,7 @@ class P2pNcclEngine:
529
comm, cudaStream_t(stream.cuda_stream))
530
stream.synchronize()
531
515
- def _recv(self, comm, tensor: torch.Tensor, src: int, stream=None):
532
+ def recv(self, comm, tensor: torch.Tensor, src: int, stream=None):
533
assert tensor.device == self.device, (
534
f"this nccl communicator is created to work on {self.device}, "
535
f"but the input tensor is on {tensor.device}")
@@ -531,3 +548,21 @@ class P2pNcclEngine:
548
self._send_thread.join()
549
if self._ping_thread is not None:
550
self._ping_thread.join()
551
+
552
+ @staticmethod
553
+ def extract_kv_from_layer(
554
+ is_mla: bool,
555
+ layer: torch.Tensor,
556
+ slot_mapping: torch.Tensor,
557
+ ) -> torch.Tensor:
558
+ """Extract the KV cache from the layer.
559
+ Assume the shape of the layer is (2, num_pages, page_size, xxx)
560
+ if MLA is not used, and (num_pages, page_size, xxx) otherwise.
561
+ """
562
+ if is_mla:
563
+ num_pages, page_size = layer.shape[0], layer.shape[1]
564
+ return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
565
+
566
+ num_pages, page_size = layer.shape[1], layer.shape[2]
567
+ return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
568
+ ...]