@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / eccdc8318c
Fork 0

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, 06:00 UTC

[V1][P/D] An native implementation of xPyD based on P2P NCCL (#18242)

Signed-off-by: Abatom <abzhonghua@gmail.com>

Zhonghua Deng committed Jun 18, 2025 at 14:32 UTC eccdc8318c50ad3e2adfaaf99d714a8dccc6ab16
8 files changed +1780
docs/design/v1/p2p_nccl_connector.md new
+337
@@ -0,0 +1,337 @@
1 +An implementation of xPyD with dynamic scaling based on point-to-point communication, partly inspired by Dynamo.
2 +
3 +# Detailed Design
4 +
5 +## Overall Process
6 +As shown in Figure 1, the overall process of this **PD disaggregation** solution is described through a request flow:
7 +
8 +1. The client sends an HTTP request to the Proxy/Router's `/v1/completions` interface.
9 +2. The Proxy/Router selects a **1P1D (1 Prefill instance + 1 Decode instance)** through either through round-robin or random selection, generates a `request_id` (rules to be introduced later), modifies the `max_tokens` in the HTTP request message to **1**, and then forwards the request to the **P instance**.
10 +3. Immediately afterward, the Proxy/Router forwards the **original HTTP request** to the **D instance**.
11 +4. The **P instance** performs **Prefill** and then **actively sends the generated KV cache** to the D instance (using **PUT_ASYNC** mode). The D instance's `zmq_addr` can be resolved through the `request_id`.
12 +5. The **D instance** has a **dedicated thread** for receiving the KV cache (to avoid blocking the main process). The received KV cache is saved into the **GPU memory buffer**, the size of which is determined by the vLLM startup parameter `kv_buffer_size`. When the GPU buffer is full, the KV cache is stored in the **local Tensor memory pool**.
13 +6. During the **Decode**, the D instance's main process retrieves the KV cache (transmitted by the P instance) from either the **GPU buffer** or the **memory pool**, thereby **skipping Prefill**.
14 +7. After completing **Decode**, the D instance returns the result to the **Proxy/Router**, which then forwards it to the **client**.
15 +
16 +![image1](https://github.com/user-attachments/assets/fb01bde6-755b-49f7-ad45-48a94b1e10a7)
17 +
18 +## Proxy/Router (Demo)
19 +
20 +A simple HTTP service acts as the entry point for client requests and starts a background thread to listen for P/D instances reporting their HTTP IP and PORT, as well as ZMQ IP and PORT. It maintains a dictionary of `http_addr -> zmq_addr`. The `http_addr` is the IP:PORT for the vLLM instance's request, while the `zmq_addr` is the address for KV cache handshake and metadata reception.
21 +
22 +The Proxy/Router is responsible for selecting 1P1D based on the characteristics of the client request, such as the prompt, and generating a corresponding `request_id`, for example:
23 +
24 +```
25 +cmpl-___prefill_addr_10.0.1.2:21001___decode_addr_10.0.1.3:22001_93923d63113b4b338973f24d19d4bf11-0
26 +```
27 +
28 +Currently, to quickly verify whether xPyD can work, a round-robin selection of 1P1D is used. In the future, it is planned to use a trie combined with the load status of instances to select appropriate P and D.
29 +
30 +Each P/D instance periodically sends a heartbeat packet to the Proxy/Router (currently every 3 seconds) to register (i.e., report `http_addr -> zmq_addr`) and keep the connection alive. If an instance crashes and fails to send a ping for a certain period of time, the Proxy/Router will remove the timed-out instance (this feature has not yet been developed).
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.
35 +
36 +Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
37 +
38 +## P2P Communication via ZMQ & NCCL
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.
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.
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.
49 +
50 +![image2](https://github.com/user-attachments/assets/837e61d6-365e-4cbf-8640-6dd7ab295b36)
51 +
52 +Each NCCL group occupies a certain amount of GPU memory buffer for communication, the size of which is primarily influenced by the `NCCL_MAX_NCHANNELS` environment variable. When `NCCL_MAX_NCHANNELS=16`, an NCCL group typically occupies 100MB, while when `NCCL_MAX_NCHANNELS=8`, it usually takes up 52MB. For large-scale xPyD configurations—such as DeepSeek's 96P144D—this implementation is currently not feasible. Moving forward, we are considering using RDMA for point-to-point communication and are also keeping an eye on UCCL.
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.
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.
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.
61 +
62 +# Install vLLM
63 +
64 +```shell
65 +# Enter the home directory or your working directory.
66 +cd /home
67 +
68 +# Download the installation package, and I will update the commit-id in time. You can directly copy the command.
69 +wget https://vllm-wheels.s3.us-west-2.amazonaws.com/9112b443a042d8d815880b8780633882ad32b183/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
70 +
71 +# Download the code repository.
72 +git clone -b xpyd-v1 https://github.com/Abatom/vllm.git
73 +cd vllm
74 +
75 +# Set the installation package path.
76 +export VLLM_PRECOMPILED_WHEEL_LOCATION=/home/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
77 +
78 +# installation
79 +pip install -e . -v
80 +```
81 +
82 +# Run xPyD
83 +
84 +## Instructions
85 +- The following examples are run on an A800 (80GB) device, using the Meta-Llama-3.1-8B-Instruct model.
86 +- Pay attention to the setting of the `kv_buffer_size` (in bytes). The empirical value is 10% of the GPU memory size. This is related to the kvcache size. If it is too small, the GPU memory buffer for temporarily storing the received kvcache will overflow, causing the kvcache to be stored in the tensor memory pool, which increases latency. If it is too large, the kvcache available for inference will be reduced, leading to a smaller batch size and decreased throughput.
87 +- For Prefill instances, when using non-GET mode, the `kv_buffer_size` can be set to 1, as Prefill currently does not need to receive kvcache. However, when using GET mode, a larger `kv_buffer_size` is required because it needs to store the kvcache sent to the D instance.
88 +- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
89 +- `PUT_ASYNC` offers the best performance and should be prioritized.
90 +- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
91 +- 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).
92 +- The node running the proxy must have `quart` installed.
93 +- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
94 +- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
95 +
96 +## Run 1P3D
97 +
98 +### Proxy (e.g. 10.0.1.1)
99 +
100 +```shell
101 +cd {your vllm directory}/examples/online_serving/disagg_xpyd/
102 +python3 disagg_prefill_proxy_xpyd.py &
103 +```
104 +
105 +### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
106 +
107 +```shell
108 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
109 + --host 0.0.0.0 \
110 + --port 20005 \
111 + --tensor-parallel-size 1 \
112 + --seed 1024 \
113 + --served-model-name base_model \
114 + --dtype float16 \
115 + --max-model-len 10000 \
116 + --max-num-batched-tokens 10000 \
117 + --max-num-seqs 256 \
118 + --trust-remote-code \
119 + --gpu-memory-utilization 0.9 \
120 + --disable-log-request \
121 + --kv-transfer-config \
122 + '{"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 &
123 +```
124 +
125 +### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
126 +
127 +```shell
128 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
129 + --host 0.0.0.0 \
130 + --port 20009 \
131 + --tensor-parallel-size 1 \
132 + --seed 1024 \
133 + --served-model-name base_model \
134 + --dtype float16 \
135 + --max-model-len 10000 \
136 + --max-num-batched-tokens 10000 \
137 + --max-num-seqs 256 \
138 + --trust-remote-code \
139 + --gpu-memory-utilization 0.7 \
140 + --disable-log-request \
141 + --kv-transfer-config \
142 + '{"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 &
143 +```
144 +
145 +### Decode2 (e.g. 10.0.1.4 or 10.0.1.1)
146 +
147 +```shell
148 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
149 + --host 0.0.0.0 \
150 + --port 20003 \
151 + --tensor-parallel-size 1 \
152 + --seed 1024 \
153 + --served-model-name base_model \
154 + --dtype float16 \
155 + --max-model-len 10000 \
156 + --max-num-batched-tokens 10000 \
157 + --max-num-seqs 256 \
158 + --trust-remote-code \
159 + --gpu-memory-utilization 0.7 \
160 + --disable-log-request \
161 + --kv-transfer-config \
162 + '{"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 &
163 +```
164 +
165 +### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
166 +
167 +```shell
168 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
169 + --host 0.0.0.0 \
170 + --port 20008 \
171 + --tensor-parallel-size 1 \
172 + --seed 1024 \
173 + --served-model-name base_model \
174 + --dtype float16 \
175 + --max-model-len 10000 \
176 + --max-num-batched-tokens 10000 \
177 + --max-num-seqs 256 \
178 + --trust-remote-code \
179 + --gpu-memory-utilization 0.7 \
180 + --disable-log-request \
181 + --kv-transfer-config \
182 + '{"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 &
183 +```
184 +
185 +## Run 3P1D
186 +
187 +### Proxy (e.g. 10.0.1.1)
188 +
189 +```shell
190 +cd {your vllm directory}/examples/online_serving/disagg_xpyd/
191 +python3 disagg_prefill_proxy_xpyd.py &
192 +```
193 +
194 +### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
195 +
196 +```shell
197 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
198 + --host 0.0.0.0 \
199 + --port 20005 \
200 + --tensor-parallel-size 1 \
201 + --seed 1024 \
202 + --served-model-name base_model \
203 + --dtype float16 \
204 + --max-model-len 10000 \
205 + --max-num-batched-tokens 10000 \
206 + --max-num-seqs 256 \
207 + --trust-remote-code \
208 + --gpu-memory-utilization 0.9 \
209 + --disable-log-request \
210 + --kv-transfer-config \
211 + '{"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 &
212 +```
213 +
214 +### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
215 +
216 +```shell
217 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
218 + --host 0.0.0.0 \
219 + --port 20009 \
220 + --tensor-parallel-size 1 \
221 + --seed 1024 \
222 + --served-model-name base_model \
223 + --dtype float16 \
224 + --max-model-len 10000 \
225 + --max-num-batched-tokens 10000 \
226 + --max-num-seqs 256 \
227 + --trust-remote-code \
228 + --gpu-memory-utilization 0.9 \
229 + --disable-log-request \
230 + --kv-transfer-config \
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":"20009","send_type":"PUT_ASYNC","nccl_num_channels":"16"}}' > /var/vllm.log 2>&1 &
232 +```
233 +
234 +### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
235 +
236 +```shell
237 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
238 + --host 0.0.0.0 \
239 + --port 20003 \
240 + --tensor-parallel-size 1 \
241 + --seed 1024 \
242 + --served-model-name base_model \
243 + --dtype float16 \
244 + --max-model-len 10000 \
245 + --max-num-batched-tokens 10000 \
246 + --max-num-seqs 256 \
247 + --trust-remote-code \
248 + --gpu-memory-utilization 0.9 \
249 + --disable-log-request \
250 + --kv-transfer-config \
251 + '{"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 &
252 +```
253 +
254 +### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
255 +
256 +```shell
257 +VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
258 + --host 0.0.0.0 \
259 + --port 20008 \
260 + --tensor-parallel-size 1 \
261 + --seed 1024 \
262 + --served-model-name base_model \
263 + --dtype float16 \
264 + --max-model-len 10000 \
265 + --max-num-batched-tokens 10000 \
266 + --max-num-seqs 256 \
267 + --trust-remote-code \
268 + --gpu-memory-utilization 0.7 \
269 + --disable-log-request \
270 + --kv-transfer-config \
271 + '{"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 &
272 +```
273 +
274 +# Single request
275 +
276 +```shell
277 +curl -X POST -s http://10.0.1.1:10001/v1/completions \
278 +-H "Content-Type: application/json" \
279 +-d '{
280 + "model": "base_model",
281 + "prompt": "San Francisco is a",
282 + "max_tokens": 10,
283 + "temperature": 0
284 +}'
285 +```
286 +
287 +# Benchmark
288 +
289 +```shell
290 +python3 benchmark_serving.py \
291 + --backend vllm \
292 + --model base_model \
293 + --tokenizer meta-llama/Llama-3.1-8B-Instruct \
294 + --dataset-name "random" \
295 + --host 10.0.1.1 \
296 + --port 10001 \
297 + --random-input-len 1024 \
298 + --random-output-len 1024 \
299 + --ignore-eos \
300 + --burstiness 100 \
301 + --percentile-metrics "ttft,tpot,itl,e2el" \
302 + --metric-percentiles "90,95,99" \
303 + --seed $(date +%s) \
304 + --trust-remote-code \
305 + --request-rate 3 \
306 + --num-prompts 1000
307 +```
308 +
309 +# Shut down
310 +
311 +```shell
312 +pgrep python | xargs kill -9 && pkill -f python
313 +```
314 +
315 +# Test data
316 +
317 +## **Scenario 1**: 1K input & 1K output tokens, E2E P99 latency ~20s
318 +- **1P5D (6×A800) vs vLLM (1×A800)**:
319 + - Throughput ↑7.2% (1085 → 6979/6)
320 + - ITL (P99) ↓81.3% (120ms → 22.9ms)
321 + - TTFT (P99) ↑26.8% (175ms → 222ms)
322 + - TPOT: No change
323 +
324 +- **1P6D (7×A800) vs vLLM (1×A800)**:
325 + - Throughput ↑9.6% (1085 → 8329/7)
326 + - ITL (P99) ↓81.0% (120ms → 22.7ms)
327 + - TTFT (P99) ↑210% (175ms →543ms)
328 + - TPOT: No change
329 +
330 +## **Scenario 2**: 1K input & 200 output tokens, E2E P99 latency ~4s
331 +- **1P1D (2×A800) vs vLLM (1×A800)**:
332 + - Throughput ↑37.4% (537 → 1476/2)
333 + - ITL (P99) ↓81.8% (127ms → 23.1ms)
334 + - TTFT (P99) ↑41.8% (160ms → 227ms)
335 + - TPOT: No change
336 +
337 +![testdata](https://github.com/user-attachments/assets/f791bfc7-9f3d-4e5c-9171-a42f9f4da627)
examples/online_serving/disagg_xpyd/disagg_prefill_proxy_xpyd.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import os
4 +import socket
5 +import threading
6 +import uuid
7 +
8 +import aiohttp
9 +import msgpack
10 +import zmq
11 +from quart import Quart, make_response, request
12 +
13 +count = 0
14 +prefill_instances: dict[str, str] = {} # http_address: zmq_address
15 +decode_instances: dict[str, str] = {} # http_address: zmq_address
16 +
17 +prefill_cv = threading.Condition()
18 +decode_cv = threading.Condition()
19 +
20 +
21 +def _listen_for_register(poller, router_socket):
22 + while True:
23 + socks = dict(poller.poll())
24 + if router_socket in socks:
25 + remote_address, message = router_socket.recv_multipart()
26 + # data: {"type": "P", "http_address": "ip:port",
27 + # "zmq_address": "ip:port"}
28 + data = msgpack.loads(message)
29 + if data["type"] == "P":
30 + global prefill_instances
31 + global prefill_cv
32 + with prefill_cv:
33 + prefill_instances[data["http_address"]] = data["zmq_address"]
34 + elif data["type"] == "D":
35 + global decode_instances
36 + global decode_cv
37 + with decode_cv:
38 + decode_instances[data["http_address"]] = data["zmq_address"]
39 + else:
40 + print(
41 + "Unexpected, Received message from %s, data: %s",
42 + remote_address,
43 + data,
44 + )
45 +
46 +
47 +def start_service_discovery(hostname, port):
48 + if not hostname:
49 + hostname = socket.gethostname()
50 + if port == 0:
51 + raise ValueError("Port cannot be 0")
52 +
53 + context = zmq.Context()
54 + router_socket = context.socket(zmq.ROUTER)
55 + router_socket.bind(f"tcp://{hostname}:{port}")
56 +
57 + poller = zmq.Poller()
58 + poller.register(router_socket, zmq.POLLIN)
59 +
60 + _listener_thread = threading.Thread(
61 + target=_listen_for_register, args=[poller, router_socket], daemon=True
62 + )
63 + _listener_thread.start()
64 + return _listener_thread
65 +
66 +
67 +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
68 +
69 +app = Quart(__name__)
70 +
71 +
72 +def random_uuid() -> str:
73 + return str(uuid.uuid4().hex)
74 +
75 +
76 +async def forward_request(url, data, request_id):
77 + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
78 + headers = {
79 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
80 + "X-Request-Id": request_id,
81 + }
82 + async with session.post(url=url, json=data, headers=headers) as response:
83 + if response.status == 200:
84 + if True:
85 + async for chunk_bytes in response.content.iter_chunked(1024):
86 + yield chunk_bytes
87 + else:
88 + content = await response.read()
89 + yield content
90 +
91 +
92 +@app.route("/v1/completions", methods=["POST"])
93 +async def handle_request():
94 + try:
95 + original_request_data = await request.get_json()
96 +
97 + prefill_request = original_request_data.copy()
98 + # change max_tokens = 1 to let it only do prefill
99 + prefill_request["max_tokens"] = 1
100 +
101 + global count
102 + global prefill_instances
103 + global prefill_cv
104 + with prefill_cv:
105 + prefill_list = list(prefill_instances.items())
106 + prefill_addr, prefill_zmq_addr = prefill_list[count % len(prefill_list)]
107 +
108 + global decode_instances
109 + global decode_cv
110 + with decode_cv:
111 + decode_list = list(decode_instances.items())
112 + decode_addr, decode_zmq_addr = decode_list[count % len(decode_list)]
113 +
114 + print(
115 + f"handle_request count: {count}, [HTTP:{prefill_addr}, "
116 + f"ZMQ:{prefill_zmq_addr}] 👉 [HTTP:{decode_addr}, "
117 + f"ZMQ:{decode_zmq_addr}]"
118 + )
119 + count += 1
120 +
121 + request_id = (
122 + f"___prefill_addr_{prefill_zmq_addr}___decode_addr_"
123 + f"{decode_zmq_addr}_{random_uuid()}"
124 + )
125 +
126 + # finish prefill
127 + async for _ in forward_request(
128 + f"http://{prefill_addr}/v1/completions", prefill_request, request_id
129 + ):
130 + continue
131 +
132 + # return decode
133 + generator = forward_request(
134 + f"http://{decode_addr}/v1/completions", original_request_data, request_id
135 + )
136 + response = await make_response(generator)
137 + response.timeout = None
138 +
139 + return response
140 +
141 + except Exception as e:
142 + import sys
143 + import traceback
144 +
145 + exc_info = sys.exc_info()
146 + print("Error occurred in disagg prefill proxy server")
147 + print(e)
148 + print("".join(traceback.format_exception(*exc_info)))
149 +
150 +
151 +if __name__ == "__main__":
152 + t = start_service_discovery("0.0.0.0", 30001)
153 + app.run(host="0.0.0.0", port=10001)
154 + t.join()
vllm/distributed/device_communicators/pynccl_wrapper.py
+8
@@ -272,6 +272,14 @@ class NCCLLibrary:
272 ctypes.byref(unique_id)))
273 return unique_id
274
275 + def unique_id_from_bytes(self, data: bytes) -> ncclUniqueId:
276 + if len(data) != 128:
277 + raise ValueError(
278 + f"Expected 128 bytes for ncclUniqueId, got {len(data)} bytes")
279 + unique_id = ncclUniqueId()
280 + ctypes.memmove(ctypes.addressof(unique_id.internal), data, 128)
281 + return unique_id
282 +
283 def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId,
284 rank: int) -> ncclComm_t:
285 comm = ncclComm_t()
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -112,6 +112,11 @@ KVConnectorFactory.register_connector(
112 "vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector",
113 "SharedStorageConnector")
114
115 +KVConnectorFactory.register_connector(
116 + "P2pNcclConnector",
117 + "vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_connector",
118 + "P2pNcclConnector")
119 +
120 KVConnectorFactory.register_connector(
121 "LMCacheConnectorV1",
122 "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",
vllm/distributed/kv_transfer/kv_connector/v1/p2p/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py new
+481
@@ -0,0 +1,481 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from dataclasses import dataclass
4 +from typing import TYPE_CHECKING, Any, Optional
5 +
6 +import regex as re
7 +import torch
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
11 + KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
12 +from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
13 + P2pNcclEngine)
14 +from vllm.distributed.parallel_state import get_world_group
15 +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
19 +
20 +if TYPE_CHECKING:
21 + from vllm.attention.backends.abstract import AttentionMetadata
22 + from vllm.forward_context import ForwardContext
23 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24 + from vllm.v1.request import Request
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +@dataclass
30 +class ReqMeta:
31 + # Request Id
32 + request_id: str
33 + # Request tokens
34 + token_ids: torch.Tensor
35 + # Slot mappings, should have the same length as token_ids
36 + slot_mapping: torch.Tensor
37 +
38 + @staticmethod
39 + def make_meta(request_id: str, token_ids: list[int], block_ids: list[int],
40 + block_size: int) -> "ReqMeta":
41 + valid_num_tokens = len(token_ids)
42 + token_ids_tensor = torch.tensor(token_ids)
43 + block_ids_tensor = torch.tensor(block_ids)
44 + num_blocks = block_ids_tensor.shape[0]
45 + block_offsets = torch.arange(0, block_size)
46 + slot_mapping = block_offsets.reshape((1, block_size)) + \
47 + block_ids_tensor.reshape((num_blocks, 1)) * block_size
48 + slot_mapping = slot_mapping.flatten()[:valid_num_tokens]
49 +
50 + return ReqMeta(
51 + request_id=request_id,
52 + token_ids=token_ids_tensor,
53 + slot_mapping=slot_mapping,
54 + )
55 +
56 +
57 +@dataclass
58 +class P2pNcclConnectorMetadata(KVConnectorMetadata):
59 + requests: list[ReqMeta]
60 +
61 + def __init__(self):
62 + self.requests = []
63 +
64 + def add_request(
65 + self,
66 + request_id: str,
67 + token_ids: list[int],
68 + block_ids: list[int],
69 + block_size: int,
70 + ) -> None:
71 + self.requests.append(
72 + ReqMeta.make_meta(request_id, token_ids, block_ids, block_size))
73 +
74 +
75 +class P2pNcclConnector(KVConnectorBase_V1):
76 +
77 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
78 + super().__init__(vllm_config=vllm_config, role=role)
79 + self._block_size = vllm_config.cache_config.block_size
80 + self._requests_need_load: dict[str, Any] = {}
81 + self.config = vllm_config.kv_transfer_config
82 + self.is_producer = self.config.is_kv_producer
83 + self.chunked_prefill: dict[str, Any] = {}
84 +
85 + self._rank = get_world_group().rank \
86 + if role == KVConnectorRole.WORKER else 0
87 + self._local_rank = get_world_group().local_rank \
88 + if role == KVConnectorRole.WORKER else 0
89 +
90 + self.p2p_nccl_engine = P2pNcclEngine(
91 + local_rank=self._local_rank,
92 + config=self.config,
93 + hostname="",
94 + port_offset=self._rank,
95 + ) if role == KVConnectorRole.WORKER else None
96 +
97 + # ==============================
98 + # Worker-side methods
99 + # ==============================
100 +
101 + def start_load_kv(self, forward_context: "ForwardContext",
102 + **kwargs) -> None:
103 + """Start loading the KV cache from the connector buffer to vLLM's
104 + paged KV buffer.
105 +
106 + Args:
107 + forward_context (ForwardContext): the forward context.
108 + **kwargs: additional arguments for the load operation
109 +
110 + Note:
111 + The number of elements in kv_caches and layer_names should be
112 + the same.
113 + """
114 +
115 + # Only consumer/decode loads KV Cache
116 + if self.is_producer:
117 + return
118 +
119 + assert self.p2p_nccl_engine is not None
120 +
121 + attn_metadata = forward_context.attn_metadata
122 + if attn_metadata is None:
123 + return
124 +
125 + def inject_kv_into_layer(
126 + dst_kv_cache_layer: torch.Tensor,
127 + src_kv_cache: torch.Tensor,
128 + slot_mapping: torch.Tensor,
129 + request_id: str,
130 + ) -> None:
131 + """Inject the KV cache into the layer.
132 +
133 + Args:
134 + dst_kv_cache_layer (torch.Tensor): the destination KV cache
135 + layer. In shape [2, num_pages, page_size, xxx] if not
136 + using MLA, [num_pages, page_size, xxx] otherwise.
137 + src_kv_cache (torch.Tensor): the source KV cache. In shape
138 + [2, num_tokens, xxx] if not using MLA, [num_tokens, xxx]
139 + otherwise.
140 + slot_mapping (torch.Tensor): the slot mapping. In shape
141 + [num_tokens].
142 + request_id (str): request id for log
143 + """
144 + dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
145 + if isinstance(attn_metadata, MLACommonMetadata):
146 + num_pages = dst_kv_cache_layer_shape[0]
147 + page_size = dst_kv_cache_layer_shape[1]
148 + dst_kv_cache_layer = dst_kv_cache_layer.reshape(
149 + num_pages * page_size, -1)
150 + self.check_tensors_except_dim(dst_kv_cache_layer, src_kv_cache,
151 + 0)
152 + num_token = src_kv_cache.shape[0]
153 + if len(slot_mapping) == num_token:
154 + dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
155 + else:
156 + dst_kv_cache_layer[slot_mapping[:num_token],
157 + ...] = src_kv_cache
158 + logger.warning(
159 + "🚧src_kv_cache does not match, num_slot:%d, "
160 + "num_token:%d, request_id:%s", len(slot_mapping),
161 + num_token, request_id)
162 +
163 + dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
164 + else:
165 + num_pages = dst_kv_cache_layer_shape[1]
166 + page_size = dst_kv_cache_layer_shape[2]
167 + dst_kv_cache_layer = dst_kv_cache_layer.reshape(
168 + 2, num_pages * page_size, -1)
169 + self.check_tensors_except_dim(dst_kv_cache_layer, src_kv_cache,
170 + 1)
171 + num_token = src_kv_cache.shape[1]
172 + if len(slot_mapping) == num_token:
173 + dst_kv_cache_layer[:, slot_mapping, ...] = src_kv_cache
174 + else:
175 + dst_kv_cache_layer[:, slot_mapping[:num_token],
176 + ...] = src_kv_cache
177 + logger.warning(
178 + "🚧src_kv_cache does not match, num_slot:%d, "
179 + "num_token:%d, request_id:%s", len(slot_mapping),
180 + num_token, request_id)
181 +
182 + dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
183 +
184 + # Get the metadata
185 + metadata: KVConnectorMetadata = \
186 + self._get_connector_metadata()
187 + assert isinstance(metadata, P2pNcclConnectorMetadata)
188 +
189 + if metadata is None:
190 + return
191 +
192 + # Load the KV for each request each layer
193 + for request in metadata.requests:
194 + for layer_name in forward_context.no_compile_layers:
195 + attn_layer = forward_context.no_compile_layers[layer_name]
196 + kv_cache_layer = attn_layer.kv_cache[ \
197 + forward_context.virtual_engine]
198 +
199 + kv_cache = self.p2p_nccl_engine.recv_tensor(
200 + request.request_id + "#" + layer_name)
201 +
202 + if kv_cache is None:
203 + logger.warning("🚧src_kv_cache is None, %s",
204 + request.request_id)
205 + continue
206 +
207 + inject_kv_into_layer(kv_cache_layer, kv_cache,
208 + request.slot_mapping, request.request_id)
209 +
210 + def wait_for_layer_load(self, layer_name: str) -> None:
211 + """Blocking until the KV for a specific layer is loaded into vLLM's
212 + paged buffer.
213 +
214 + This interface will be useful for layer-by-layer pipelining.
215 +
216 + Args:
217 + layer_name: the name of that layer
218 + """
219 + return
220 +
221 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
222 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
223 + """Start saving the KV cache of the layer from vLLM's paged buffer
224 + to the connector.
225 +
226 + Args:
227 + layer_name (str): the name of the layer.
228 + kv_layer (torch.Tensor): the paged KV buffer of the current
229 + layer in vLLM.
230 + attn_metadata (AttentionMetadata): the attention metadata.
231 + **kwargs: additional arguments for the save operation.
232 + """
233 +
234 + # Only producer/prefill saves KV Cache
235 + if not self.is_producer:
236 + return
237 +
238 + assert self.p2p_nccl_engine is not None
239 +
240 + def extract_kv_from_layer(
241 + layer: torch.Tensor,
242 + slot_mapping: torch.Tensor,
243 + ) -> torch.Tensor:
244 + """Extract the KV cache from the layer.
245 +
246 + Assume the shape of the layer is (2, num_pages, page_size, xxx)
247 + if MLA is not used, and (num_pages, page_size, xxx) otherwise.
248 + """
249 + if isinstance(attn_metadata, MLACommonMetadata):
250 + num_pages, page_size = layer.shape[0], layer.shape[1]
251 + return layer.reshape(num_pages * page_size, -1)[slot_mapping,
252 + ...]
253 + num_pages, page_size = layer.shape[1], layer.shape[2]
254 + return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
255 + ...]
256 +
257 + connector_metadata = self._get_connector_metadata()
258 + assert isinstance(connector_metadata, P2pNcclConnectorMetadata)
259 + for request in connector_metadata.requests:
260 + request_id = request.request_id
261 + ip, port = self.parse_request_id(request_id, True)
262 + remote_address = ip + ":" + str(port + self._rank)
263 + kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
264 + self.p2p_nccl_engine.send_tensor(request_id + "#" + layer_name,
265 + kv_cache, remote_address)
266 +
267 + def wait_for_save(self):
268 + if self.is_producer:
269 + assert self.p2p_nccl_engine is not None
270 + self.p2p_nccl_engine.wait_for_sent()
271 +
272 + def get_finished(
273 + self, finished_req_ids: set[str],
274 + **kwargs) -> tuple[Optional[set[str]], Optional[set[str]]]:
275 + """
276 + Notifies worker-side connector ids of requests that have
277 + finished generating tokens.
278 +
279 + Returns:
280 + ids of requests that have finished asynchronous transfer,
281 + tuple of (sending/saving ids, recving/loading ids).
282 + The finished saves/sends req ids must belong to a set provided in a
283 + call to this method (this call or a prior one).
284 + """
285 +
286 + assert self.p2p_nccl_engine is not None
287 +
288 + forward_context: ForwardContext = get_forward_context()
289 + return self.p2p_nccl_engine.get_finished(finished_req_ids,
290 + forward_context)
291 +
292 + # ==============================
293 + # Scheduler-side methods
294 + # ==============================
295 +
296 + def get_num_new_matched_tokens(
297 + self,
298 + request: "Request",
299 + num_computed_tokens: int,
300 + ) -> tuple[int, bool]:
301 + """
302 + Get number of new tokens that can be loaded from the
303 + external KV cache beyond the num_computed_tokens.
304 +
305 + Args:
306 + request (Request): the request object.
307 + num_computed_tokens (int): the number of locally
308 + computed tokens for this request
309 +
310 + Returns:
311 + the number of tokens that can be loaded from the
312 + external KV cache beyond what is already computed.
313 + """
314 + if self.is_producer:
315 + return 0, False
316 +
317 + num_external_tokens = (len(request.prompt_token_ids) - 1 -
318 + num_computed_tokens)
319 +
320 + if num_external_tokens < 0:
321 + num_external_tokens = 0
322 +
323 + return num_external_tokens, False
324 +
325 + def update_state_after_alloc(self, request: "Request",
326 + blocks: "KVCacheBlocks",
327 + num_external_tokens: int):
328 + """
329 + Update KVConnector state after block allocation.
330 + """
331 + if not self.is_producer and num_external_tokens > 0:
332 + self._requests_need_load[request.request_id] = (
333 + request, blocks.get_block_ids()[0])
334 +
335 + def build_connector_meta(
336 + self,
337 + scheduler_output: SchedulerOutput,
338 + ) -> KVConnectorMetadata:
339 + """Build the connector metadata for this step.
340 +
341 + This function should NOT modify any fields in the scheduler_output.
342 + Also, calling this function will reset the state of the connector.
343 +
344 + Args:
345 + scheduler_output (SchedulerOutput): the scheduler output object.
346 + """
347 +
348 + meta = P2pNcclConnectorMetadata()
349 +
350 + for new_req in scheduler_output.scheduled_new_reqs:
351 + if self.is_producer:
352 + num_scheduled_tokens = (
353 + scheduler_output.num_scheduled_tokens)[new_req.req_id]
354 + num_tokens = num_scheduled_tokens + new_req.num_computed_tokens
355 + # the request's prompt is chunked prefill
356 + if num_tokens < len(new_req.prompt_token_ids):
357 + # 'CachedRequestData' has no attribute 'prompt_token_ids'
358 + self.chunked_prefill[new_req.req_id] = (
359 + new_req.block_ids[0], new_req.prompt_token_ids)
360 + continue
361 + # the request's prompt is not chunked prefill
362 + meta.add_request(request_id=new_req.req_id,
363 + token_ids=new_req.prompt_token_ids,
364 + block_ids=new_req.block_ids[0],
365 + block_size=self._block_size)
366 + continue
367 + if new_req.req_id in self._requests_need_load:
368 + meta.add_request(request_id=new_req.req_id,
369 + token_ids=new_req.prompt_token_ids,
370 + block_ids=new_req.block_ids[0],
371 + block_size=self._block_size)
372 + self._requests_need_load.pop(new_req.req_id)
373 +
374 + for cached_req in scheduler_output.scheduled_cached_reqs:
375 + if self.is_producer:
376 + num_scheduled_tokens = (
377 + scheduler_output.num_scheduled_tokens)[cached_req.req_id]
378 + num_tokens = (num_scheduled_tokens +
379 + cached_req.num_computed_tokens)
380 + assert cached_req.req_id in self.chunked_prefill
381 + block_ids = cached_req.new_block_ids[0]
382 + if not cached_req.resumed_from_preemption:
383 + block_ids = (self.chunked_prefill[cached_req.req_id][0] +
384 + block_ids)
385 + prompt_token_ids = self.chunked_prefill[cached_req.req_id][1]
386 + # the request's prompt is chunked prefill again
387 + if num_tokens < len(prompt_token_ids):
388 + self.chunked_prefill[cached_req.req_id] = (
389 + block_ids, prompt_token_ids)
390 + continue
391 + # the request's prompt is all prefilled finally
392 + meta.add_request(request_id=cached_req.req_id,
393 + token_ids=prompt_token_ids,
394 + block_ids=block_ids,
395 + block_size=self._block_size)
396 + self.chunked_prefill.pop(cached_req.req_id, None)
397 + continue
398 +
399 + # NOTE(rob): here we rely on the resumed requests being
400 + # the first N requests in the list scheduled_cache_reqs.
401 + if not cached_req.resumed_from_preemption:
402 + break
403 + if cached_req.req_id in self._requests_need_load:
404 + request, _ = self._requests_need_load.pop(cached_req.req_id)
405 + total_tokens = cached_req.num_computed_tokens + 1
406 + token_ids = request.all_token_ids[:total_tokens]
407 +
408 + # NOTE(rob): For resumed req, new_block_ids is all
409 + # of the block_ids for the request.
410 + block_ids = cached_req.new_block_ids[0]
411 +
412 + meta.add_request(request_id=cached_req.req_id,
413 + token_ids=token_ids,
414 + block_ids=block_ids,
415 + block_size=self._block_size)
416 +
417 + # Requests loaded asynchronously are not in the scheduler_output.
418 + # for request_id in self._requests_need_load:
419 + # request, block_ids = self._requests_need_load[request_id]
420 + # meta.add_request(request_id=request.request_id,
421 + # token_ids=request.prompt_token_ids,
422 + # block_ids=block_ids,
423 + # block_size=self._block_size)
424 +
425 + self._requests_need_load.clear()
426 + return meta
427 +
428 + def request_finished(
429 + self,
430 + request: "Request",
431 + block_ids: list[int],
432 + ) -> tuple[bool, Optional[dict[str, Any]]]:
433 + """
434 + Called when a request has finished, before its blocks are freed.
435 +
436 + Returns:
437 + True if the request is being saved/sent asynchronously and blocks
438 + should not be freed until the request_id is returned from
439 + get_finished().
440 + Optional KVTransferParams to be included in the request outputs
441 + returned by the engine.
442 + """
443 +
444 + self.chunked_prefill.pop(request.request_id, None)
445 +
446 + return False, None
447 +
448 + # ==============================
449 + # Static methods
450 + # ==============================
451 +
452 + @staticmethod
453 + def parse_request_id(request_id: str, is_prefill=True) -> tuple[str, int]:
454 + # Regular expression to match the string hostname and integer port
455 + if is_prefill:
456 + pattern = r"___decode_addr_(.*):(\d+)"
457 + else:
458 + pattern = r"___prefill_addr_(.*):(\d+)___"
459 +
460 + # Use re.search to find the pattern in the request_id
461 + match = re.search(pattern, request_id)
462 + if match:
463 + # Extract the ranks
464 + ip = match.group(1)
465 + port = int(match.group(2))
466 +
467 + return ip, port
468 + raise ValueError(
469 + f"Request id {request_id} does not contain hostname and port")
470 +
471 + @staticmethod
472 + def check_tensors_except_dim(tensor1, tensor2, dim):
473 + shape1 = tensor1.size()
474 + shape2 = tensor2.size()
475 +
476 + if len(shape1) != len(shape2) or not all(
477 + s1 == s2
478 + for i, (s1, s2) in enumerate(zip(shape1, shape2)) if i != dim):
479 + raise NotImplementedError(
480 + "Currently, only symmetric TP is supported. Asymmetric TP, PP,"
481 + "and others will be supported in future PRs.")
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py new
+531
@@ -0,0 +1,531 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import logging
4 +import os
5 +import threading
6 +import time
7 +import typing
8 +from collections import deque
9 +from contextlib import contextmanager
10 +from typing import TYPE_CHECKING, Any, Optional
11 +
12 +import msgpack
13 +import torch
14 +import zmq
15 +
16 +from vllm.config import KVTransferConfig
17 +from vllm.distributed.device_communicators.pynccl_wrapper import (
18 + NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum)
19 +from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import ( # noqa: E501
20 + TensorMemoryPool)
21 +from vllm.utils import current_stream, get_ip
22 +
23 +if TYPE_CHECKING:
24 + from vllm.forward_context import ForwardContext
25 +
26 +logger = logging.getLogger(__name__)
27 +
28 +DEFAULT_MEM_POOL_SIZE_GB = 32
29 +
30 +
31 +@contextmanager
32 +def set_p2p_nccl_context(num_channels: str):
33 + original_values: dict[str, Any] = {}
34 + env_vars = [
35 + 'NCCL_MAX_NCHANNELS',
36 + 'NCCL_MIN_NCHANNELS',
37 + 'NCCL_CUMEM_ENABLE',
38 + 'NCCL_BUFFSIZE',
39 + 'NCCL_PROTO', # LL,LL128,SIMPLE
40 + 'NCCL_ALGO', # RING,TREE
41 + ]
42 +
43 + for var in env_vars:
44 + original_values[var] = os.environ.get(var)
45 +
46 + logger.info("set_p2p_nccl_context, original_values: %s", original_values)
47 +
48 + try:
49 + os.environ['NCCL_MAX_NCHANNELS'] = num_channels
50 + os.environ['NCCL_MIN_NCHANNELS'] = num_channels
51 + os.environ['NCCL_CUMEM_ENABLE'] = '1'
52 + yield
53 + finally:
54 + for var in env_vars:
55 + if original_values[var] is not None:
56 + os.environ[var] = original_values[var]
57 + else:
58 + os.environ.pop(var, None)
59 +
60 +
61 +class P2pNcclEngine:
62 +
63 + def __init__(self,
64 + local_rank: int,
65 + config: KVTransferConfig,
66 + hostname: str = "",
67 + port_offset: int = 0,
68 + library_path: Optional[str] = None) -> None:
69 + self.config = config
70 + self.rank = port_offset
71 + self.local_rank = local_rank
72 + self.device = torch.device(f"cuda:{self.local_rank}")
73 + self.nccl = NCCLLibrary(library_path)
74 +
75 + if not hostname:
76 + hostname = get_ip()
77 + port = int(self.config.kv_port) + port_offset
78 + if port == 0:
79 + raise ValueError("Port cannot be 0")
80 + self._hostname = hostname
81 + self._port = port
82 +
83 + # Each card corresponds to a ZMQ address.
84 + self.zmq_address = f"{self._hostname}:{self._port}"
85 +
86 + # The `http_port` must be consistent with the port of OpenAI.
87 + self.http_address = (
88 + f"{self._hostname}:"
89 + f"{self.config.kv_connector_extra_config['http_port']}")
90 +
91 + # If `proxy_ip` or `proxy_port` is `""`,
92 + # then the ping thread will not be enabled.
93 + proxy_ip = self.config.get_from_extra_config("proxy_ip", "")
94 + proxy_port = self.config.get_from_extra_config("proxy_port", "")
95 + if proxy_ip == "" or proxy_port == "":
96 + self.proxy_address = ""
97 + else:
98 + self.proxy_address = proxy_ip + ":" + proxy_port
99 +
100 + self.context = zmq.Context()
101 + self.router_socket = self.context.socket(zmq.ROUTER)
102 + self.router_socket.bind(f"tcp://{self.zmq_address}")
103 +
104 + self.poller = zmq.Poller()
105 + self.poller.register(self.router_socket, zmq.POLLIN)
106 +
107 + self.send_store_cv = threading.Condition()
108 + self.send_queue_cv = threading.Condition()
109 + self.recv_store_cv = threading.Condition()
110 +
111 + self.send_stream = torch.cuda.Stream()
112 + self.recv_stream = torch.cuda.Stream()
113 +
114 + mem_pool_size_gb = self.config.get_from_extra_config(
115 + "mem_pool_size_gb", DEFAULT_MEM_POOL_SIZE_GB)
116 + self.pool = TensorMemoryPool(max_block_size=int(mem_pool_size_gb) *
117 + 1024**3) # GB
118 +
119 + # The sending type includes tree mutually exclusive options:
120 + # PUT, GET, PUT_ASYNC.
121 + self.send_type = self.config.get_from_extra_config("send_type", "PUT")
122 + if self.send_type == "GET":
123 + # tensor_id: torch.Tensor
124 + self.send_store: dict[str, torch.Tensor] = {}
125 + else:
126 + # PUT or PUT_ASYNC
127 + # tensor_id: torch.Tensor
128 + self.send_queue: deque[list[Any]] = deque()
129 + self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
130 + if self.send_type == "PUT_ASYNC":
131 + self._send_thread = threading.Thread(target=self._send_async,
132 + daemon=True)
133 + self._send_thread.start()
134 +
135 + # tensor_id: torch.Tensor/(addr, dtype, shape)
136 + self.recv_store: dict[str, Any] = {}
137 + self.recv_request_id_to_tensor_ids: dict[str, set[str]] = {}
138 + self.socks: dict[str, Any] = {} # remote_address: client socket
139 + self.comms: dict[str, Any] = {} # remote_address: (ncclComm_t, rank)
140 +
141 + self.buffer_size = 0
142 + self.buffer_size_threshold = float(self.config.kv_buffer_size)
143 +
144 + self.nccl_num_channels = self.config.get_from_extra_config(
145 + "nccl_num_channels", "8")
146 +
147 + self._listener_thread = threading.Thread(
148 + target=self._listen_for_requests, daemon=True)
149 + self._listener_thread.start()
150 +
151 + self._ping_thread = None
152 + if port_offset == 0 and self.proxy_address != "":
153 + self._ping_thread = threading.Thread(target=self._ping,
154 + daemon=True)
155 + self._ping_thread.start()
156 +
157 + logger.info(
158 + "💯P2pNcclEngine init, rank:%d, local_rank:%d, http_address:%s, "
159 + "zmq_address:%s, proxy_address:%s, send_type:%s, buffer_size_"
160 + "threshold:%.2f, nccl_num_channels:%s", self.rank, self.local_rank,
161 + self.http_address, self.zmq_address, self.proxy_address,
162 + self.send_type, self.buffer_size_threshold, self.nccl_num_channels)
163 +
164 + def _create_connect(self, remote_address: typing.Optional[str] = None):
165 + assert remote_address is not None
166 + if remote_address not in self.socks:
167 + sock = self.context.socket(zmq.DEALER)
168 + sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
169 + sock.connect(f"tcp://{remote_address}")
170 + self.socks[remote_address] = sock
171 + if remote_address in self.comms:
172 + logger.info("👋comm exists, remote_address:%s, comms:%s",
173 + remote_address, self.comms)
174 + return sock, self.comms[remote_address]
175 +
176 + unique_id = self.nccl.ncclGetUniqueId()
177 + data = {"cmd": "NEW", "unique_id": bytes(unique_id.internal)}
178 + sock.send(msgpack.dumps(data))
179 +
180 + with torch.cuda.device(self.device):
181 + rank = 0
182 + with set_p2p_nccl_context(self.nccl_num_channels):
183 + comm: ncclComm_t = self.nccl.ncclCommInitRank(
184 + 2, unique_id, rank)
185 + self.comms[remote_address] = (comm, rank)
186 + logger.info("🤝ncclCommInitRank Success, %s👉%s, MyRank: %s",
187 + self.zmq_address, remote_address, rank)
188 +
189 + return self.socks[remote_address], self.comms[remote_address]
190 +
191 + def send_tensor(
192 + self,
193 + tensor_id: str,
194 + tensor: torch.Tensor,
195 + remote_address: typing.Optional[str] = None,
196 + ) -> bool:
197 + if remote_address is None:
198 + with self.recv_store_cv:
199 + self.recv_store[tensor_id] = tensor
200 + self.recv_store_cv.notify()
201 + return True
202 + else:
203 + if self.send_type == "PUT":
204 + return self._send_sync(tensor_id, tensor, remote_address)
205 + elif self.send_type == "PUT_ASYNC":
206 + with self.send_queue_cv:
207 + self.send_queue.append([tensor_id, remote_address, tensor])
208 + self.send_queue_cv.notify()
209 + else: # GET
210 + with self.send_store_cv:
211 + tensor_size = tensor.element_size() * tensor.numel()
212 + while (self.buffer_size + tensor_size
213 + > self.buffer_size_threshold):
214 + oldest_tenser_id = next(iter(self.send_store))
215 + oldest_tenser = self.send_store.pop(oldest_tenser_id)
216 + oldest_tenser_size = oldest_tenser.element_size(
217 + ) * oldest_tenser.numel()
218 + self.buffer_size -= oldest_tenser_size
219 + logger.info(
220 + "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
221 + " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
222 + remote_address, tensor_id, tensor_size,
223 + self.buffer_size, oldest_tenser_size, self.rank)
224 +
225 + self.send_store[tensor_id] = tensor
226 + self.buffer_size += tensor_size
227 + logger.debug(
228 + "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
229 + "shape:%s, rank:%d, buffer_size:%d(%.2f%%)",
230 + remote_address, tensor_id, tensor_size, tensor.shape,
231 + self.rank, self.buffer_size,
232 + self.buffer_size / self.buffer_size_threshold * 100)
233 +
234 + return True
235 +
236 + def recv_tensor(
237 + self,
238 + tensor_id: str,
239 + remote_address: typing.Optional[str] = None,
240 + ) -> torch.Tensor:
241 + if self.send_type == "PUT" or self.send_type == "PUT_ASYNC":
242 + start_time = time.time()
243 + with self.recv_store_cv:
244 + while tensor_id not in self.recv_store:
245 + self.recv_store_cv.wait()
246 + tensor = self.recv_store[tensor_id]
247 +
248 + if tensor is not None:
249 + if isinstance(tensor, tuple):
250 + addr, dtype, shape = tensor
251 + tensor = self.pool.load_tensor(addr, dtype, shape,
252 + self.device)
253 + else:
254 + self.buffer_size -= (tensor.element_size() *
255 + tensor.numel())
256 + else:
257 + duration = time.time() - start_time
258 + logger.warning(
259 + "🔴[PUT]Recv From %s, tensor_id:%s, duration:%.3fms, "
260 + "rank:%d", remote_address, tensor_id, duration * 1000,
261 + self.rank)
262 + return tensor
263 +
264 + # GET
265 + if remote_address is None:
266 + return None
267 +
268 + if remote_address not in self.socks:
269 + self._create_connect(remote_address)
270 +
271 + sock = self.socks[remote_address]
272 + comm, rank = self.comms[remote_address]
273 +
274 + data = {"cmd": "GET", "tensor_id": tensor_id}
275 + sock.send(msgpack.dumps(data))
276 +
277 + message = sock.recv()
278 + data = msgpack.loads(message)
279 + if data["ret"] != 0:
280 + logger.warning("🔴[GET]Recv From %s, tensor_id: %s, ret: %d",
281 + remote_address, tensor_id, data["ret"])
282 + return None
283 +
284 + tensor = torch.empty(data["shape"],
285 + dtype=getattr(torch, data["dtype"]),
286 + device=self.device)
287 +
288 + self._recv(comm, tensor, rank ^ 1, self.recv_stream)
289 +
290 + return tensor
291 +
292 + def _listen_for_requests(self):
293 + while True:
294 + socks = dict(self.poller.poll())
295 + if self.router_socket in socks:
296 + remote_address, message = self.router_socket.recv_multipart()
297 + data = msgpack.loads(message)
298 + if data["cmd"] == "NEW":
299 + unique_id = self.nccl.unique_id_from_bytes(
300 + bytes(data["unique_id"]))
301 + with torch.cuda.device(self.device):
302 + rank = 1
303 + with set_p2p_nccl_context(self.nccl_num_channels):
304 + comm: ncclComm_t = self.nccl.ncclCommInitRank(
305 + 2, unique_id, rank)
306 + self.comms[remote_address.decode()] = (comm, rank)
307 + logger.info(
308 + "🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
309 + self.zmq_address, remote_address.decode(), rank)
310 + elif data["cmd"] == "PUT":
311 + tensor_id = data["tensor_id"]
312 + try:
313 + tensor = torch.empty(data["shape"],
314 + dtype=getattr(
315 + torch, data["dtype"]),
316 + device=self.device)
317 + self.router_socket.send_multipart(
318 + [remote_address, b"0"])
319 + comm, rank = self.comms[remote_address.decode()]
320 + self._recv(comm, tensor, rank ^ 1, self.recv_stream)
321 + tensor_size = tensor.element_size() * tensor.numel()
322 + if (self.buffer_size + tensor_size
323 + > self.buffer_size_threshold):
324 + # Store Tensor in memory pool
325 + addr = self.pool.store_tensor(tensor)
326 + tensor = (addr, tensor.dtype, tensor.shape)
327 + logger.warning(
328 + "🔴[PUT]Recv Tensor, Out Of Threshold, "
329 + "%s👈%s, data:%s, addr:%d", self.zmq_address,
330 + remote_address.decode(), data, addr)
331 + else:
332 + self.buffer_size += tensor_size
333 +
334 + except torch.cuda.OutOfMemoryError:
335 + self.router_socket.send_multipart(
336 + [remote_address, b"1"])
337 + tensor = None
338 + logger.warning(
339 + "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, "
340 + "data:%s", self.zmq_address,
341 + remote_address.decode(), data)
342 +
343 + with self.recv_store_cv:
344 + self.recv_store[tensor_id] = tensor
345 + self._have_received_tensor_id(tensor_id)
346 + self.recv_store_cv.notify()
347 +
348 + elif data["cmd"] == "GET":
349 + tensor_id = data["tensor_id"]
350 + with self.send_store_cv:
351 + tensor = self.send_store.pop(tensor_id, None)
352 + if tensor is not None:
353 + data = {
354 + "ret": 0,
355 + "shape": tensor.shape,
356 + "dtype":
357 + str(tensor.dtype).replace("torch.", "")
358 + }
359 + # LRU
360 + self.send_store[tensor_id] = tensor
361 + self._have_sent_tensor_id(tensor_id)
362 + else:
363 + data = {"ret": 1}
364 +
365 + self.router_socket.send_multipart(
366 + [remote_address, msgpack.dumps(data)])
367 +
368 + if data["ret"] == 0:
369 + comm, rank = self.comms[remote_address.decode()]
370 + self._send(comm, tensor.to(self.device), rank ^ 1,
371 + self.send_stream)
372 + else:
373 + logger.warning(
374 + "🚧Unexpected, Received message from %s, data:%s",
375 + remote_address, data)
376 +
377 + def _have_sent_tensor_id(self, tensor_id: str):
378 + request_id = tensor_id.split('#')[0]
379 + if request_id not in self.send_request_id_to_tensor_ids:
380 + self.send_request_id_to_tensor_ids[request_id] = set()
381 + self.send_request_id_to_tensor_ids[request_id].add(tensor_id)
382 +
383 + def _have_received_tensor_id(self, tensor_id: str):
384 + request_id = tensor_id.split('#')[0]
385 + if request_id not in self.recv_request_id_to_tensor_ids:
386 + self.recv_request_id_to_tensor_ids[request_id] = set()
387 + self.recv_request_id_to_tensor_ids[request_id].add(tensor_id)
388 +
389 + def _send_async(self):
390 + while True:
391 + with self.send_queue_cv:
392 + while not self.send_queue:
393 + self.send_queue_cv.wait()
394 + tensor_id, remote_address, tensor = self.send_queue.popleft()
395 + if not self.send_queue:
396 + self.send_queue_cv.notify()
397 + self._send_sync(tensor_id, tensor, remote_address)
398 +
399 + def wait_for_sent(self):
400 + if self.send_type == "PUT_ASYNC":
401 + start_time = time.time()
402 + with self.send_queue_cv:
403 + while self.send_queue:
404 + self.send_queue_cv.wait()
405 + duration = time.time() - start_time
406 + logger.debug(
407 + "🚧[PUT_ASYNC]It took %.3fms to wait for the send_queue"
408 + " to be empty, rank:%d", duration * 1000, self.rank)
409 +
410 + def _send_sync(
411 + self,
412 + tensor_id: str,
413 + tensor: torch.Tensor,
414 + remote_address: typing.Optional[str] = None,
415 + ) -> bool:
416 + if remote_address is None:
417 + return False
418 + if remote_address not in self.socks:
419 + self._create_connect(remote_address)
420 +
421 + sock = self.socks[remote_address]
422 + comm, rank = self.comms[remote_address]
423 + data = {
424 + "cmd": "PUT",
425 + "tensor_id": tensor_id,
426 + "shape": tensor.shape,
427 + "dtype": str(tensor.dtype).replace("torch.", "")
428 + }
429 + sock.send(msgpack.dumps(data))
430 +
431 + response = sock.recv()
432 + if response != b"0":
433 + logger.error(
434 + "🔴Send Tensor, Peer Out Of Memory/Threshold, %s 👉 %s, "
435 + "MyRank:%s, data:%s, tensor:%s, size:%fGB, response:%s",
436 + self.zmq_address, remote_address, rank, data, tensor.shape,
437 + tensor.element_size() * tensor.numel() / 1024**3,
438 + response.decode())
439 + return False
440 +
441 + self._send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
442 +
443 + if self.send_type == "PUT_ASYNC":
444 + self._have_sent_tensor_id(tensor_id)
445 +
446 + return True
447 +
448 + def get_finished(
449 + self, finished_req_ids: set[str], forward_context: "ForwardContext"
450 + ) -> tuple[Optional[set[str]], Optional[set[str]]]:
451 + """
452 + Notifies worker-side connector ids of requests that have
453 + finished generating tokens.
454 +
455 + Returns:
456 + ids of requests that have finished asynchronous transfer,
457 + tuple of (sending/saving ids, recving/loading ids).
458 + The finished saves/sends req ids must belong to a set provided in a
459 + call to this method (this call or a prior one).
460 + """
461 +
462 + # Clear the buffer upon request completion.
463 + for request_id in finished_req_ids:
464 + for layer_name in forward_context.no_compile_layers:
465 + tensor_id = request_id + "#" + layer_name
466 + if tensor_id in self.recv_store:
467 + with self.recv_store_cv:
468 + tensor = self.recv_store.pop(tensor_id, None)
469 + self.send_request_id_to_tensor_ids.pop(
470 + request_id, None)
471 + self.recv_request_id_to_tensor_ids.pop(
472 + request_id, None)
473 + addr = 0
474 + if isinstance(tensor, tuple):
475 + addr, _, _ = tensor
476 + self.pool.free(addr)
477 +
478 + # TODO:Retrieve requests that have already sent the KV cache.
479 + finished_sending: set[str] = set()
480 +
481 + # TODO:Retrieve requests that have already received the KV cache.
482 + finished_recving: set[str] = set()
483 +
484 + return finished_sending or None, finished_recving or None
485 +
486 + def _ping(self):
487 + sock = self.context.socket(zmq.DEALER)
488 + sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
489 + logger.debug("ping start, zmq_address:%s", self.zmq_address)
490 + sock.connect(f"tcp://{self.proxy_address}")
491 + data = {
492 + "type": "P" if self.config.is_kv_producer else "D",
493 + "http_address": self.http_address,
494 + "zmq_address": self.zmq_address
495 + }
496 + while True:
497 + sock.send(msgpack.dumps(data))
498 + time.sleep(3)
499 +
500 + def _send(self, comm, tensor: torch.Tensor, dst: int, stream=None):
501 + assert tensor.device == self.device, (
502 + f"this nccl communicator is created to work on {self.device}, "
503 + f"but the input tensor is on {tensor.device}")
504 + if stream is None:
505 + stream = current_stream()
506 +
507 + with torch.cuda.stream(stream):
508 + self.nccl.ncclSend(buffer_type(tensor.data_ptr()), tensor.numel(),
509 + ncclDataTypeEnum.from_torch(tensor.dtype), dst,
510 + comm, cudaStream_t(stream.cuda_stream))
511 + stream.synchronize()
512 +
513 + def _recv(self, comm, tensor: torch.Tensor, src: int, stream=None):
514 + assert tensor.device == self.device, (
515 + f"this nccl communicator is created to work on {self.device}, "
516 + f"but the input tensor is on {tensor.device}")
517 + if stream is None:
518 + stream = current_stream()
519 +
520 + with torch.cuda.stream(stream):
521 + self.nccl.ncclRecv(buffer_type(tensor.data_ptr()), tensor.numel(),
522 + ncclDataTypeEnum.from_torch(tensor.dtype), src,
523 + comm, cudaStream_t(stream.cuda_stream))
524 + stream.synchronize()
525 +
526 + def close(self) -> None:
527 + self._listener_thread.join()
528 + if self.send_type == "PUT_ASYNC":
529 + self._send_thread.join()
530 + if self._ping_thread is not None:
531 + self._ping_thread.join()
vllm/distributed/kv_transfer/kv_connector/v1/p2p/tensor_memory_pool.py new
+264
@@ -0,0 +1,264 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import atexit
4 +import ctypes
5 +import math
6 +from dataclasses import dataclass
7 +
8 +import torch
9 +
10 +from vllm.logger import init_logger
11 +
12 +logger = init_logger(__name__)
13 +
14 +
15 +@dataclass
16 +class MemoryBlock:
17 + size: int
18 + addr: int
19 +
20 +
21 +"""A memory pool for managing pinned host memory allocations for tensors.
22 +
23 +This class implements a buddy allocation system to efficiently manage pinned
24 +host memory for tensor storage. It supports allocation, deallocation, and
25 +tensor storage/retrieval operations.
26 +
27 +Key Features:
28 +- Uses power-of-two block sizes for efficient buddy allocation
29 +- Supports splitting and merging of memory blocks
30 +- Provides methods to store CUDA tensors in pinned host memory
31 +- Allows loading tensors from pinned memory back to device
32 +- Automatically cleans up memory on destruction
33 +
34 +Attributes:
35 + max_block_size (int): Maximum block size (rounded to nearest power of two)
36 + min_block_size (int): Minimum block size (rounded to nearest power of two)
37 + free_lists (dict): Dictionary of free memory blocks by size
38 + allocated_blocks (dict): Dictionary of currently allocated blocks
39 + base_tensor (torch.Tensor): Base pinned memory tensor
40 + base_address (int): Base memory address of the pinned memory region
41 +
42 +Example:
43 + >>> pool = TensorMemoryPool(max_block_size=1024*1024)
44 + >>> tensor = torch.randn(100, device='cuda')
45 + >>> addr = pool.store_tensor(tensor)
46 + >>> loaded_tensor = pool.load_tensor(addr, tensor.dtype,
47 + ... tensor.shape, 'cuda')
48 + >>> pool.free(addr)
49 +"""
50 +
51 +
52 +class TensorMemoryPool:
53 + """Initializes the memory pool with given size constraints.
54 +
55 + Args:
56 + max_block_size (int): Maximum size of memory blocks to manage
57 + min_block_size (int, optional): Minimum size of memory blocks
58 + to manage. Defaults to 512.
59 +
60 + Raises:
61 + ValueError: If block sizes are invalid or max_block_size is less
62 + than min_block_size
63 + """
64 +
65 + def __init__(self, max_block_size: int, min_block_size: int = 512):
66 + if max_block_size <= 0 or min_block_size <= 0:
67 + raise ValueError("Block sizes must be positive")
68 + if max_block_size < min_block_size:
69 + raise ValueError(
70 + "Max block size must be greater than min block size")
71 +
72 + self.max_block_size = self._round_to_power_of_two(max_block_size)
73 + self.min_block_size = self._round_to_power_of_two(min_block_size)
74 +
75 + self.free_lists: dict[int, dict[int, MemoryBlock]] = {}
76 + self.allocated_blocks: dict[int, MemoryBlock] = {}
77 +
78 + self._initialize_free_lists()
79 + self._allocate_pinned_memory()
80 +
81 + atexit.register(self.cleanup)
82 +
83 + def _round_to_power_of_two(self, size: int) -> int:
84 + return 1 << (size - 1).bit_length()
85 +
86 + def _initialize_free_lists(self):
87 + size = self.max_block_size
88 + while size >= self.min_block_size:
89 + self.free_lists[size] = {}
90 + size //= 2
91 +
92 + def _allocate_pinned_memory(self):
93 + self.base_tensor = torch.empty(self.max_block_size // 4,
94 + dtype=torch.float32,
95 + pin_memory=True)
96 + self.base_address = self.base_tensor.data_ptr()
97 + initial_block = MemoryBlock(size=self.max_block_size,
98 + addr=self.base_address)
99 + self.free_lists[self.max_block_size][
100 + initial_block.addr] = initial_block
101 + logger.debug("TensorMemoryPool, base_address:", self.base_address,
102 + self.base_address % self.max_block_size)
103 +
104 + def allocate(self, size: int) -> int:
105 + """Allocates a memory block of at least the requested size.
106 +
107 + Args:
108 + size (int): Minimum size of memory to allocate
109 +
110 + Returns:
111 + int: Address of the allocated memory block
112 +
113 + Raises:
114 + ValueError: If size is invalid or insufficient memory is available
115 + """
116 + if size <= 0:
117 + raise ValueError("Allocation size must be positive")
118 +
119 + required_size = self._round_to_power_of_two(
120 + max(size, self.min_block_size))
121 + if required_size > self.max_block_size:
122 + raise ValueError("Requested size exceeds maximum block size")
123 +
124 + current_size = required_size
125 + while current_size <= self.max_block_size:
126 + if self.free_lists[current_size]:
127 + _, block = self.free_lists[current_size].popitem()
128 + self._split_block(block, required_size)
129 + self.allocated_blocks[block.addr] = block
130 + return block.addr
131 + current_size *= 2
132 +
133 + raise ValueError("Insufficient memory")
134 +
135 + def _split_block(self, block: MemoryBlock, required_size: int):
136 + while (block.size > required_size
137 + and block.size // 2 >= self.min_block_size):
138 + buddy_size = block.size // 2
139 + buddy_addr = block.addr + buddy_size
140 +
141 + buddy = MemoryBlock(size=buddy_size, addr=buddy_addr)
142 + block.size = buddy_size
143 +
144 + self.free_lists[buddy_size][buddy.addr] = buddy
145 +
146 + def free(self, addr: int):
147 + """Frees an allocated memory block.
148 +
149 + Args:
150 + addr (int): Address of the block to free
151 +
152 + Raises:
153 + ValueError: If address is invalid or not allocated
154 + """
155 + if addr not in self.allocated_blocks:
156 + raise ValueError("Invalid address to free")
157 +
158 + block = self.allocated_blocks.pop(addr)
159 + self._merge_buddies(block)
160 +
161 + def _merge_buddies(self, block: MemoryBlock):
162 + MAX_MERGE_DEPTH = 30
163 + depth = 0
164 +
165 + while depth < MAX_MERGE_DEPTH:
166 + buddy_offset = block.size if (block.addr - self.base_address) % (
167 + 2 * block.size) == 0 else -block.size
168 + buddy_addr = block.addr + buddy_offset
169 + buddy = self.free_lists[block.size].get(buddy_addr)
170 + if buddy:
171 + del self.free_lists[buddy.size][buddy.addr]
172 + merged_addr = min(block.addr, buddy.addr)
173 + merged_size = block.size * 2
174 + block = MemoryBlock(size=merged_size, addr=merged_addr)
175 + depth += 1
176 + else:
177 + break
178 + self.free_lists[block.size][block.addr] = block
179 +
180 + def store_tensor(self, tensor: torch.Tensor) -> int:
181 + """Stores a CUDA tensor in pinned host memory.
182 +
183 + Args:
184 + tensor (torch.Tensor): CUDA tensor to store
185 +
186 + Returns:
187 + int: Address where the tensor is stored
188 +
189 + Raises:
190 + ValueError: If tensor is not on CUDA or allocation fails
191 + """
192 + if not tensor.is_cuda:
193 + raise ValueError("Only CUDA tensors can be stored")
194 +
195 + size = tensor.element_size() * tensor.numel()
196 + addr = self.allocate(size)
197 + block = self.allocated_blocks[addr]
198 +
199 + if block.size < size:
200 + self.free(addr)
201 + raise ValueError(
202 + f"Allocated block size {block.size} is smaller than "
203 + f"required size {size}")
204 +
205 + try:
206 + buffer = (ctypes.c_byte * block.size).from_address(block.addr)
207 + cpu_tensor = torch.frombuffer(buffer,
208 + dtype=tensor.dtype,
209 + count=tensor.numel()).reshape(
210 + tensor.shape)
211 + except ValueError as err:
212 + self.free(addr)
213 + raise ValueError(f"Failed to create tensor view: {err}") from err
214 +
215 + cpu_tensor.copy_(tensor)
216 +
217 + return addr
218 +
219 + def load_tensor(self, addr: int, dtype: torch.dtype,
220 + shape: tuple[int, ...], device) -> torch.Tensor:
221 + """Loads a tensor from pinned host memory to the specified device.
222 +
223 + Args:
224 + addr (int): Address where tensor is stored
225 + dtype (torch.dtype): Data type of the tensor
226 + shape (tuple[int, ...]): Shape of the tensor
227 + device: Target device for the loaded tensor
228 +
229 + Returns:
230 + torch.Tensor: The loaded tensor on the specified device
231 +
232 + Raises:
233 + ValueError: If address is invalid or sizes don't match
234 + """
235 + if addr not in self.allocated_blocks:
236 + raise ValueError("Invalid address to load")
237 +
238 + block = self.allocated_blocks[addr]
239 + num_elements = math.prod(shape)
240 + dtype_size = torch.tensor([], dtype=dtype).element_size()
241 + required_size = num_elements * dtype_size
242 +
243 + if required_size > block.size:
244 + raise ValueError("Requested tensor size exceeds block size")
245 +
246 + buffer = (ctypes.c_byte * block.size).from_address(block.addr)
247 + cpu_tensor = torch.frombuffer(buffer, dtype=dtype,
248 + count=num_elements).reshape(shape)
249 +
250 + cuda_tensor = torch.empty(shape, dtype=dtype, device=device)
251 +
252 + cuda_tensor.copy_(cpu_tensor)
253 +
254 + return cuda_tensor
255 +
256 + def cleanup(self):
257 + """Cleans up all memory resources and resets the pool state."""
258 + self.free_lists.clear()
259 + self.allocated_blocks.clear()
260 + if hasattr(self, 'base_tensor'):
261 + del self.base_tensor
262 +
263 + def __del__(self):
264 + self.cleanup()