A high-throughput and memory-efficient inference and serving engine for LLMs
[Connector] Remove `P2pNcclConnector` (#44854)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 8, 2026 at 12:58 UTC
5add018bebbd54471af32094923205fe4c9b7740
18 files changed
+2
-3069
benchmarks/disagg_benchmarks/disagg_overhead_benchmark.sh
deleted
-143
@@ -1,143 +0,0 @@
1
-#!/bin/bash
2
-
3
-# benchmark the overhead of disaggregated prefill.
4
-# methodology:
5
-# - send all request to prefill vLLM instance. It will buffer KV cache.
6
-# - then send all request to decode instance.
7
-# - The TTFT of decode instance is the overhead.
8
-
9
-set -ex
10
-
11
-kill_gpu_processes() {
12
- # kill all processes on GPU.
13
- pgrep pt_main_thread | xargs -r kill -9
14
- pgrep python3 | xargs -r kill -9
15
- # vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
16
- pgrep VLLM | xargs -r kill -9
17
- sleep 10
18
-
19
- # remove vllm config file
20
- rm -rf ~/.config/vllm
21
-
22
- # Print the GPU memory usage
23
- # so that we know if all GPU processes are killed.
24
- gpu_memory_usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0)
25
- # The memory usage should be 0 MB.
26
- echo "GPU 0 Memory Usage: $gpu_memory_usage MB"
27
-}
28
-
29
-wait_for_server() {
30
- # wait for vllm server to start
31
- # return 1 if vllm server crashes
32
- local port=$1
33
- timeout 1200 bash -c "
34
- until curl -s localhost:${port}/v1/completions > /dev/null; do
35
- sleep 1
36
- done" && return 0 || return 1
37
-}
38
-
39
-
40
-benchmark() {
41
-
42
- export VLLM_LOGGING_LEVEL=DEBUG
43
- export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
44
-
45
- # compare chunked prefill with disaggregated prefill
46
-
47
- results_folder="./results"
48
- model="meta-llama/Meta-Llama-3.1-8B-Instruct"
49
- dataset_name="sonnet"
50
- dataset_path="../sonnet_4x.txt"
51
- num_prompts=10
52
- qps=$1
53
- prefix_len=50
54
- input_len=2048
55
- output_len=$2
56
-
57
-
58
- CUDA_VISIBLE_DEVICES=0 vllm serve $model \
59
- --port 8100 \
60
- --max-model-len 10000 \
61
- --gpu-memory-utilization 0.6 \
62
- --kv-transfer-config \
63
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
64
-
65
-
66
- CUDA_VISIBLE_DEVICES=1 vllm serve $model \
67
- --port 8200 \
68
- --max-model-len 10000 \
69
- --gpu-memory-utilization 0.6 \
70
- --kv-transfer-config \
71
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
72
-
73
- wait_for_server 8100
74
- wait_for_server 8200
75
-
76
- # let the prefill instance finish prefill
77
- vllm bench serve \
78
- --backend vllm \
79
- --model $model \
80
- --dataset-name $dataset_name \
81
- --dataset-path $dataset_path \
82
- --sonnet-input-len $input_len \
83
- --sonnet-output-len "$output_len" \
84
- --sonnet-prefix-len $prefix_len \
85
- --num-prompts $num_prompts \
86
- --port 8100 \
87
- --save-result \
88
- --result-dir $results_folder \
89
- --result-filename disagg_prefill_tp1.json \
90
- --request-rate "inf"
91
-
92
-
93
- # send the request to decode.
94
- # The TTFT of this command will be the overhead of disagg prefill impl.
95
- vllm bench serve \
96
- --backend vllm \
97
- --model $model \
98
- --dataset-name $dataset_name \
99
- --dataset-path $dataset_path \
100
- --sonnet-input-len $input_len \
101
- --sonnet-output-len "$output_len" \
102
- --sonnet-prefix-len $prefix_len \
103
- --num-prompts $num_prompts \
104
- --port 8200 \
105
- --save-result \
106
- --result-dir $results_folder \
107
- --result-filename disagg_prefill_tp1_overhead.json \
108
- --request-rate "$qps"
109
- kill_gpu_processes
110
-
111
-}
112
-
113
-
114
-main() {
115
-
116
- (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
117
- (which jq) || (apt-get -y install jq)
118
- (which socat) || (apt-get -y install socat)
119
-
120
- pip install quart httpx datasets
121
-
122
- cd "$(dirname "$0")"
123
-
124
- cd ..
125
- # create sonnet-4x.txt
126
- echo "" > sonnet_4x.txt
127
- for _ in {1..4}
128
- do
129
- cat sonnet.txt >> sonnet_4x.txt
130
- done
131
- cd disagg_benchmarks
132
-
133
- rm -rf results
134
- mkdir results
135
-
136
- default_qps=1
137
- default_output_len=1
138
- benchmark $default_qps $default_output_len
139
-
140
-}
141
-
142
-
143
-main "$@"
benchmarks/disagg_benchmarks/disagg_performance_benchmark.sh
deleted
-157
@@ -1,157 +0,0 @@
1
-#!/bin/bash
2
-
3
-# Requirement: 2x GPUs.
4
-
5
-
6
-# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
7
-# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
8
-# Resource: 2x GPU
9
-# Approaches:
10
-# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
11
-# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
12
-# Prefilling instance: max_output_token=1
13
-# Decoding instance: force the input tokens be the same across requests to bypass prefilling
14
-
15
-set -ex
16
-
17
-kill_gpu_processes() {
18
- # kill all processes on GPU.
19
- pgrep pt_main_thread | xargs -r kill -9
20
- pgrep python3 | xargs -r kill -9
21
- # vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
22
- pgrep VLLM | xargs -r kill -9
23
- for port in 8000 8100 8200; do lsof -t -i:$port | xargs -r kill -9; done
24
- sleep 1
25
-}
26
-
27
-wait_for_server() {
28
- # wait for vllm server to start
29
- # return 1 if vllm server crashes
30
- local port=$1
31
- timeout 1200 bash -c "
32
- until curl -s localhost:${port}/v1/completions > /dev/null; do
33
- sleep 1
34
- done" && return 0 || return 1
35
-}
36
-
37
-
38
-launch_chunked_prefill() {
39
- model="meta-llama/Meta-Llama-3.1-8B-Instruct"
40
- # disagg prefill
41
- CUDA_VISIBLE_DEVICES=0 vllm serve $model \
42
- --port 8100 \
43
- --max-model-len 10000 \
44
- --enable-chunked-prefill \
45
- --gpu-memory-utilization 0.6 &
46
- CUDA_VISIBLE_DEVICES=1 vllm serve $model \
47
- --port 8200 \
48
- --max-model-len 10000 \
49
- --enable-chunked-prefill \
50
- --gpu-memory-utilization 0.6 &
51
- wait_for_server 8100
52
- wait_for_server 8200
53
- python3 round_robin_proxy.py &
54
- sleep 1
55
-}
56
-
57
-
58
-launch_disagg_prefill() {
59
- model="meta-llama/Meta-Llama-3.1-8B-Instruct"
60
- # disagg prefill
61
- CUDA_VISIBLE_DEVICES=0 vllm serve $model \
62
- --port 8100 \
63
- --max-model-len 10000 \
64
- --gpu-memory-utilization 0.6 \
65
- --kv-transfer-config \
66
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
67
-
68
- CUDA_VISIBLE_DEVICES=1 vllm serve $model \
69
- --port 8200 \
70
- --max-model-len 10000 \
71
- --gpu-memory-utilization 0.6 \
72
- --kv-transfer-config \
73
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
74
-
75
- wait_for_server 8100
76
- wait_for_server 8200
77
- python3 disagg_prefill_proxy_server.py &
78
- sleep 1
79
-}
80
-
81
-
82
-benchmark() {
83
- results_folder="./results"
84
- model="meta-llama/Meta-Llama-3.1-8B-Instruct"
85
- dataset_name="sonnet"
86
- dataset_path="../sonnet_4x.txt"
87
- num_prompts=100
88
- qps=$1
89
- prefix_len=50
90
- input_len=1024
91
- output_len=$2
92
- tag=$3
93
-
94
- vllm bench serve \
95
- --backend vllm \
96
- --model $model \
97
- --dataset-name $dataset_name \
98
- --dataset-path $dataset_path \
99
- --sonnet-input-len $input_len \
100
- --sonnet-output-len "$output_len" \
101
- --sonnet-prefix-len $prefix_len \
102
- --num-prompts $num_prompts \
103
- --port 8000 \
104
- --save-result \
105
- --result-dir $results_folder \
106
- --result-filename "$tag"-qps-"$qps".json \
107
- --request-rate "$qps"
108
-
109
- sleep 2
110
-}
111
-
112
-
113
-main() {
114
-
115
- (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
116
- (which jq) || (apt-get -y install jq)
117
- (which socat) || (apt-get -y install socat)
118
- (which lsof) || (apt-get -y install lsof)
119
-
120
- pip install quart httpx matplotlib aiohttp datasets
121
-
122
- cd "$(dirname "$0")"
123
-
124
- cd ..
125
- # create sonnet-4x.txt so that we can sample 2048 tokens for input
126
- echo "" > sonnet_4x.txt
127
- for _ in {1..4}
128
- do
129
- cat sonnet.txt >> sonnet_4x.txt
130
- done
131
- cd disagg_benchmarks
132
-
133
- rm -rf results
134
- mkdir results
135
-
136
- default_output_len=6
137
-
138
- export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
139
-
140
- launch_chunked_prefill
141
- for qps in 2 4 6 8; do
142
- benchmark $qps $default_output_len chunked_prefill
143
- done
144
- kill_gpu_processes
145
-
146
- launch_disagg_prefill
147
- for qps in 2 4 6 8; do
148
- benchmark $qps $default_output_len disagg_prefill
149
- done
150
- kill_gpu_processes
151
-
152
- python3 visualize_benchmark_results.py
153
-
154
-}
155
-
156
-
157
-main "$@"
benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py
deleted
-260
@@ -1,260 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import argparse
5
-import asyncio
6
-import logging
7
-import os
8
-import time
9
-import uuid
10
-from urllib.parse import urlparse
11
-
12
-import aiohttp
13
-from quart import Quart, Response, make_response, request
14
-
15
-# Configure logging
16
-logging.basicConfig(level=logging.INFO)
17
-logger = logging.getLogger(__name__)
18
-
19
-
20
-def parse_args():
21
- """parse command line arguments"""
22
- parser = argparse.ArgumentParser(description="vLLM P/D disaggregation proxy server")
23
-
24
- # Add args
25
- parser.add_argument(
26
- "--timeout",
27
- type=float,
28
- default=6 * 60 * 60,
29
- help="Timeout for backend service requests in seconds (default: 21600)",
30
- )
31
- parser.add_argument(
32
- "--port",
33
- type=int,
34
- default=8000,
35
- help="Port to run the server on (default: 8000)",
36
- )
37
- parser.add_argument(
38
- "--prefill-url",
39
- type=str,
40
- default="http://localhost:8100",
41
- help="Prefill service base URL (protocol + host[:port])",
42
- )
43
- parser.add_argument(
44
- "--decode-url",
45
- type=str,
46
- default="http://localhost:8200",
47
- help="Decode service base URL (protocol + host[:port])",
48
- )
49
- parser.add_argument(
50
- "--kv-host",
51
- type=str,
52
- default="localhost",
53
- help="Hostname or IP used by KV transfer (default: localhost)",
54
- )
55
- parser.add_argument(
56
- "--prefill-kv-port",
57
- type=int,
58
- default=14579,
59
- help="Prefill KV port (default: 14579)",
60
- )
61
- parser.add_argument(
62
- "--decode-kv-port",
63
- type=int,
64
- default=14580,
65
- help="Decode KV port (default: 14580)",
66
- )
67
-
68
- return parser.parse_args()
69
-
70
-
71
-def main():
72
- """parse command line arguments"""
73
- args = parse_args()
74
-
75
- # Initialize configuration using command line parameters
76
- AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
77
- PREFILL_SERVICE_URL = args.prefill_url
78
- DECODE_SERVICE_URL = args.decode_url
79
- PORT = args.port
80
-
81
- PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
82
- DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
83
-
84
- logger.info(
85
- "Proxy resolved KV addresses -> prefill: %s, decode: %s",
86
- PREFILL_KV_ADDR,
87
- DECODE_KV_ADDR,
88
- )
89
-
90
- app = Quart(__name__)
91
-
92
- # Attach the configuration object to the application instance so helper
93
- # coroutines can read the resolved backend URLs and timeouts without using
94
- # globals.
95
- app.config.update(
96
- {
97
- "AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
98
- "PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
99
- "DECODE_SERVICE_URL": DECODE_SERVICE_URL,
100
- "PREFILL_KV_ADDR": PREFILL_KV_ADDR,
101
- "DECODE_KV_ADDR": DECODE_KV_ADDR,
102
- }
103
- )
104
-
105
- def _normalize_base_url(url: str) -> str:
106
- """Remove any trailing slash so path joins behave predictably."""
107
- return url.rstrip("/")
108
-
109
- def _get_host_port(url: str) -> str:
110
- """Return the hostname:port portion for logging and KV headers."""
111
- parsed = urlparse(url)
112
- host = parsed.hostname or "localhost"
113
- port = parsed.port
114
- if port is None:
115
- port = 80 if parsed.scheme == "http" else 443
116
- return f"{host}:{port}"
117
-
118
- PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
119
- DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
120
- KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
121
-
122
- def _build_headers(request_id: str) -> dict[str, str]:
123
- """Construct the headers expected by vLLM's P2P disagg connector."""
124
- headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
125
- api_key = os.environ.get("OPENAI_API_KEY")
126
- if api_key:
127
- headers["Authorization"] = f"Bearer {api_key}"
128
- return headers
129
-
130
- async def _run_prefill(
131
- request_path: str,
132
- payload: dict,
133
- headers: dict[str, str],
134
- request_id: str,
135
- ):
136
- url = f"{PREFILL_BASE}{request_path}"
137
- start_ts = time.perf_counter()
138
- logger.info("[prefill] start request_id=%s url=%s", request_id, url)
139
- try:
140
- async with (
141
- aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
142
- session.post(url=url, json=payload, headers=headers) as resp,
143
- ):
144
- if resp.status != 200:
145
- error_text = await resp.text()
146
- raise RuntimeError(
147
- f"Prefill backend error {resp.status}: {error_text}"
148
- )
149
- await resp.read()
150
- logger.info(
151
- "[prefill] done request_id=%s status=%s elapsed=%.2fs",
152
- request_id,
153
- resp.status,
154
- time.perf_counter() - start_ts,
155
- )
156
- except asyncio.TimeoutError as exc:
157
- raise RuntimeError(f"Prefill service timeout at {url}") from exc
158
- except aiohttp.ClientError as exc:
159
- raise RuntimeError(f"Prefill service unavailable at {url}") from exc
160
-
161
- async def _stream_decode(
162
- request_path: str,
163
- payload: dict,
164
- headers: dict[str, str],
165
- request_id: str,
166
- ):
167
- url = f"{DECODE_BASE}{request_path}"
168
- # Stream tokens from the decode service once the prefill stage has
169
- # materialized KV caches on the target workers.
170
- logger.info("[decode] start request_id=%s url=%s", request_id, url)
171
- try:
172
- async with (
173
- aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
174
- session.post(url=url, json=payload, headers=headers) as resp,
175
- ):
176
- if resp.status != 200:
177
- error_text = await resp.text()
178
- logger.error(
179
- "Decode backend error %s - %s", resp.status, error_text
180
- )
181
- err_msg = (
182
- '{"error": "Decode backend error ' + str(resp.status) + '"}'
183
- )
184
- yield err_msg.encode()
185
- return
186
- logger.info(
187
- "[decode] streaming response request_id=%s status=%s",
188
- request_id,
189
- resp.status,
190
- )
191
- async for chunk_bytes in resp.content.iter_chunked(1024):
192
- yield chunk_bytes
193
- logger.info("[decode] finished streaming request_id=%s", request_id)
194
- except asyncio.TimeoutError:
195
- logger.error("Decode service timeout at %s", url)
196
- yield b'{"error": "Decode service timeout"}'
197
- except aiohttp.ClientError as exc:
198
- logger.error("Decode service error at %s: %s", url, exc)
199
- yield b'{"error": "Decode service unavailable"}'
200
-
201
- async def process_request():
202
- """Process a single request through prefill and decode stages"""
203
- try:
204
- original_request_data = await request.get_json()
205
-
206
- # Create prefill request (max_tokens=1)
207
- prefill_request = original_request_data.copy()
208
- prefill_request["max_tokens"] = 1
209
- if "max_completion_tokens" in prefill_request:
210
- prefill_request["max_completion_tokens"] = 1
211
-
212
- # Execute prefill stage
213
- # The request id encodes both KV socket addresses so the backend can
214
- # shuttle tensors directly via NCCL once the prefill response
215
- # completes.
216
- request_id = (
217
- f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
218
- f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
219
- )
220
-
221
- headers = _build_headers(request_id)
222
- await _run_prefill(request.path, prefill_request, headers, request_id)
223
-
224
- # Execute decode stage and stream response
225
- # Pass the unmodified user request so the decode phase can continue
226
- # sampling with the already-populated KV cache.
227
- generator = _stream_decode(
228
- request.path, original_request_data, headers, request_id
229
- )
230
- response = await make_response(generator)
231
- response.timeout = None # Disable timeout for streaming response
232
- return response
233
-
234
- except Exception:
235
- logger.exception("Error processing request")
236
- return Response(
237
- response=b'{"error": "Internal server error"}',
238
- status=500,
239
- content_type="application/json",
240
- )
241
-
242
- @app.route("/v1/completions", methods=["POST"])
243
- async def handle_request():
244
- """Handle incoming API requests with concurrency and rate limiting"""
245
- try:
246
- return await process_request()
247
- except asyncio.CancelledError:
248
- logger.warning("Request cancelled")
249
- return Response(
250
- response=b'{"error": "Request cancelled"}',
251
- status=503,
252
- content_type="application/json",
253
- )
254
-
255
- # Start the Quart server with host can be set to 0.0.0.0
256
- app.run(port=PORT)
257
-
258
-
259
-if __name__ == "__main__":
260
- main()
benchmarks/disagg_benchmarks/round_robin_proxy.py
deleted
-63
@@ -1,63 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import asyncio
5
-import itertools
6
-
7
-import aiohttp
8
-from aiohttp import web
9
-
10
-
11
-class RoundRobinProxy:
12
- def __init__(self, target_ports):
13
- self.target_ports = target_ports
14
- self.port_cycle = itertools.cycle(self.target_ports)
15
-
16
- async def handle_request(self, request):
17
- target_port = next(self.port_cycle)
18
- target_url = f"http://localhost:{target_port}{request.path_qs}"
19
-
20
- async with aiohttp.ClientSession() as session:
21
- try:
22
- # Forward the request
23
- async with session.request(
24
- method=request.method,
25
- url=target_url,
26
- headers=request.headers,
27
- data=request.content,
28
- ) as response:
29
- # Start sending the response
30
- resp = web.StreamResponse(
31
- status=response.status, headers=response.headers
32
- )
33
- await resp.prepare(request)
34
-
35
- # Stream the response content
36
- async for chunk in response.content.iter_any():
37
- await resp.write(chunk)
38
-
39
- await resp.write_eof()
40
- return resp
41
-
42
- except Exception as e:
43
- return web.Response(text=f"Error: {str(e)}", status=500)
44
-
45
-
46
-async def main():
47
- proxy = RoundRobinProxy([8100, 8200])
48
- app = web.Application()
49
- app.router.add_route("*", "/{path:.*}", proxy.handle_request)
50
-
51
- runner = web.AppRunner(app)
52
- await runner.setup()
53
- site = web.TCPSite(runner, "localhost", 8000)
54
- await site.start()
55
-
56
- print("Proxy server started on http://localhost:8000")
57
-
58
- # Keep the server running
59
- await asyncio.Event().wait()
60
-
61
-
62
-if __name__ == "__main__":
63
- asyncio.run(main())
benchmarks/disagg_benchmarks/visualize_benchmark_results.py
deleted
-47
@@ -1,47 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import json
5
-
6
-import matplotlib.pyplot as plt
7
-import pandas as pd
8
-
9
-if __name__ == "__main__":
10
- data = []
11
- for name in ["disagg_prefill", "chunked_prefill"]:
12
- for qps in [2, 4, 6, 8]:
13
- with open(f"results/{name}-qps-{qps}.json") as f:
14
- x = json.load(f)
15
- x["name"] = name
16
- x["qps"] = qps
17
- data.append(x)
18
-
19
- df = pd.DataFrame.from_dict(data)
20
- dis_df = df[df["name"] == "disagg_prefill"]
21
- chu_df = df[df["name"] == "chunked_prefill"]
22
-
23
- plt.style.use("bmh")
24
- plt.rcParams["font.size"] = 20
25
-
26
- for key in [
27
- "mean_ttft_ms",
28
- "median_ttft_ms",
29
- "p99_ttft_ms",
30
- "mean_itl_ms",
31
- "median_itl_ms",
32
- "p99_itl_ms",
33
- ]:
34
- fig, ax = plt.subplots(figsize=(11, 7))
35
- plt.plot(
36
- dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4
37
- )
38
- plt.plot(
39
- chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4
40
- )
41
- ax.legend()
42
-
43
- ax.set_xlabel("QPS")
44
- ax.set_ylabel(key)
45
- ax.set_ylim(bottom=0)
46
- fig.savefig(f"results/{key}.png")
47
- plt.close(fig)
docs/design/p2p_nccl_connector.md
deleted
-319
@@ -1,319 +0,0 @@
1
-# P2P NCCL Connector
2
-
3
-An implementation of xPyD with dynamic scaling based on point-to-point communication, partly inspired by Dynamo.
4
-
5
-## Detailed Design
6
-
7
-### Overall Process
8
-
9
-As shown in Figure 1, the overall process of this **PD disaggregation** solution is described through a request flow:
10
-
11
-1. The client sends an HTTP request to the Proxy/Router's `/v1/completions` interface.
12
-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**.
13
-3. Immediately afterward, the Proxy/Router forwards the **original HTTP request** to the **D instance**.
14
-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`.
15
-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**.
16
-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**.
17
-7. After completing **Decode**, the D instance returns the result to the **Proxy/Router**, which then forwards it to the **client**.
18
-
19
-
20
-
21
-### Proxy/Router (Demo)
22
-
23
-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.
24
-
25
-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:
26
-
27
-```text
28
-cmpl-___prefill_addr_10.0.1.2:21001___decode_addr_10.0.1.3:22001_93923d63113b4b338973f24d19d4bf11-0
29
-```
30
-
31
-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.
32
-
33
-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).
34
-
35
-### KV Cache Transfer Methods
36
-
37
-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.
38
-
39
-Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
40
-
41
-### P2P Communication via ZMQ & NCCL
42
-
43
-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.
44
-
45
-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.
46
-
47
-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.
48
-
49
-### NCCL Group Topology
50
-
51
-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.
52
-
53
-
54
-
55
-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.
56
-
57
-### GPU Memory Buffer and Tensor Memory Pool
58
-
59
-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.
60
-
61
-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.
62
-
63
-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.
64
-
65
-## Install vLLM
66
-
67
-```shell
68
-pip install "vllm>=0.9.2"
69
-```
70
-
71
-## Run xPyD
72
-
73
-### Instructions
74
-
75
-- The following examples are run on an A800 (80GB) device, using the Meta-Llama-3.1-8B-Instruct model.
76
-- 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.
77
-- 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.
78
-- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
79
-- `PUT_ASYNC` offers the best performance and should be prioritized.
80
-- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
81
-- 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).
82
-- The node running the proxy must have `quart` installed.
83
-- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
84
-- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
85
-
86
-### Run 1P3D
87
-
88
-#### Proxy (e.g. 10.0.1.1)
89
-
90
-```shell
91
-cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
92
-python3 disagg_proxy_p2p_nccl_xpyd.py &
93
-```
94
-
95
-#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
96
-
97
-??? console "Command"
98
-
99
- ```shell
100
- CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
101
- --host 0.0.0.0 \
102
- --port 20001 \
103
- --tensor-parallel-size 1 \
104
- --seed 1024 \
105
- --served-model-name base_model \
106
- --dtype float16 \
107
- --max-model-len 10000 \
108
- --max-num-batched-tokens 10000 \
109
- --max-num-seqs 256 \
110
- --trust-remote-code \
111
- --gpu-memory-utilization 0.9 \
112
- --kv-transfer-config \
113
- '{"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 &
114
- ```
115
-
116
-#### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
117
-
118
-??? console "Command"
119
-
120
- ```shell
121
- CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
122
- --host 0.0.0.0 \
123
- --port 20002 \
124
- --tensor-parallel-size 1 \
125
- --seed 1024 \
126
- --served-model-name base_model \
127
- --dtype float16 \
128
- --max-model-len 10000 \
129
- --max-num-batched-tokens 10000 \
130
- --max-num-seqs 256 \
131
- --trust-remote-code \
132
- --gpu-memory-utilization 0.7 \
133
- --kv-transfer-config \
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)
138
-
139
-??? console "Command"
140
-
141
- ```shell
142
- CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
143
- --host 0.0.0.0 \
144
- --port 20003 \
145
- --tensor-parallel-size 1 \
146
- --seed 1024 \
147
- --served-model-name base_model \
148
- --dtype float16 \
149
- --max-model-len 10000 \
150
- --max-num-batched-tokens 10000 \
151
- --max-num-seqs 256 \
152
- --trust-remote-code \
153
- --gpu-memory-utilization 0.7 \
154
- --kv-transfer-config \
155
- '{"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 &
156
- ```
157
-
158
-#### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
159
-
160
-??? console "Command"
161
-
162
- ```shell
163
- CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
164
- --host 0.0.0.0 \
165
- --port 20004 \
166
- --tensor-parallel-size 1 \
167
- --seed 1024 \
168
- --served-model-name base_model \
169
- --dtype float16 \
170
- --max-model-len 10000 \
171
- --max-num-batched-tokens 10000 \
172
- --max-num-seqs 256 \
173
- --trust-remote-code \
174
- --gpu-memory-utilization 0.7 \
175
- --kv-transfer-config \
176
- '{"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 &
177
- ```
178
-
179
-### Run 3P1D
180
-
181
-#### Proxy (e.g. 10.0.1.1)
182
-
183
-```shell
184
-cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
185
-python3 disagg_proxy_p2p_nccl_xpyd.py &
186
-```
187
-
188
-#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
189
-
190
-??? console "Command"
191
-
192
- ```shell
193
- CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
194
- --host 0.0.0.0 \
195
- --port 20001 \
196
- --tensor-parallel-size 1 \
197
- --seed 1024 \
198
- --served-model-name base_model \
199
- --dtype float16 \
200
- --max-model-len 10000 \
201
- --max-num-batched-tokens 10000 \
202
- --max-num-seqs 256 \
203
- --trust-remote-code \
204
- --gpu-memory-utilization 0.9 \
205
- --kv-transfer-config \
206
- '{"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 &
207
- ```
208
-
209
-#### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
210
-
211
-??? console "Command"
212
-
213
- ```shell
214
- CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
215
- --host 0.0.0.0 \
216
- --port 20002 \
217
- --tensor-parallel-size 1 \
218
- --seed 1024 \
219
- --served-model-name base_model \
220
- --dtype float16 \
221
- --max-model-len 10000 \
222
- --max-num-batched-tokens 10000 \
223
- --max-num-seqs 256 \
224
- --trust-remote-code \
225
- --gpu-memory-utilization 0.9 \
226
- --kv-transfer-config \
227
- '{"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 &
228
- ```
229
-
230
-#### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
231
-
232
-??? console "Command"
233
-
234
- ```shell
235
- CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
236
- --host 0.0.0.0 \
237
- --port 20003 \
238
- --tensor-parallel-size 1 \
239
- --seed 1024 \
240
- --served-model-name base_model \
241
- --dtype float16 \
242
- --max-model-len 10000 \
243
- --max-num-batched-tokens 10000 \
244
- --max-num-seqs 256 \
245
- --trust-remote-code \
246
- --gpu-memory-utilization 0.9 \
247
- --kv-transfer-config \
248
- '{"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 &
249
- ```
250
-
251
-#### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
252
-
253
-??? console "Command"
254
-
255
- ```shell
256
- CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
257
- --host 0.0.0.0 \
258
- --port 20004 \
259
- --tensor-parallel-size 1 \
260
- --seed 1024 \
261
- --served-model-name base_model \
262
- --dtype float16 \
263
- --max-model-len 10000 \
264
- --max-num-batched-tokens 10000 \
265
- --max-num-seqs 256 \
266
- --trust-remote-code \
267
- --gpu-memory-utilization 0.7 \
268
- --kv-transfer-config \
269
- '{"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 &
270
- ```
271
-
272
-## Single request
273
-
274
-```shell
275
-curl -X POST -s http://10.0.1.1:10001/v1/completions \
276
--H "Content-Type: application/json" \
277
--d '{
278
- "model": "base_model",
279
- "prompt": "San Francisco is a",
280
- "max_tokens": 10,
281
- "temperature": 0
282
-}'
283
-```
284
-
285
-## Benchmark
286
-
287
-??? console "Command"
288
-
289
- ```shell
290
- vllm bench serve \
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**: 1K input & 200 output tokens, E2E P99 latency ~2s
318
-
319
-
docs/features/disagg_prefill.md
-1
@@ -29,7 +29,6 @@ Now supports 9 types of connectors:
29
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
30
```
31
32
-- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
32
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
33
- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
34
- **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
examples/disaggregated/disaggregated_prefill.py
deleted
-127
@@ -1,127 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-This file demonstrates the example usage of disaggregated prefilling
5
-We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
6
-and then transfer the KV cache between them.
7
-"""
8
-
9
-import os
10
-import time
11
-from multiprocessing import Event, Process
12
-
13
-from vllm import LLM, SamplingParams
14
-from vllm.config import KVTransferConfig
15
-
16
-
17
-def run_prefill(prefill_done):
18
- # We use GPU 0 for prefill node.
19
- os.environ["CUDA_VISIBLE_DEVICES"] = "0"
20
-
21
- # The prefill node receives two requests, while the decode node receives
22
- # three requests. So the decode node will only receive the KV Cache for
23
- # requests 1 and 3. The decode node will use the KV Cache of requests 1
24
- # and 3 and do prefilling on request 2.
25
- prompts = [
26
- "Hello, my name is",
27
- "Hi, your name is",
28
- # The decode node will actually "prefill" this request.
29
- "Tell me a very long story",
30
- ]
31
- sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
32
-
33
- # Using P2pNcclConnector to transmit KV caches between vLLM instances.
34
- # This instance is the prefill node (kv_producer, rank 0).
35
- # The number of parallel instances for KV cache transfer is set to 2,
36
- # as required for P2pNcclConnector.
37
- ktc = KVTransferConfig(
38
- kv_connector="P2pNcclConnector",
39
- kv_role="kv_producer",
40
- kv_rank=0,
41
- kv_parallel_size=2,
42
- )
43
-
44
- # Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
45
- # memory. You may need to adjust the value to fit your GPU.
46
- llm = LLM(
47
- model="meta-llama/Meta-Llama-3.1-8B-Instruct",
48
- kv_transfer_config=ktc,
49
- max_model_len=2000,
50
- gpu_memory_utilization=0.8,
51
- )
52
-
53
- llm.generate(prompts, sampling_params)
54
- print("Prefill node is finished.")
55
- prefill_done.set()
56
-
57
- # To keep the prefill node running in case the decode node is not done;
58
- # otherwise, the script might exit prematurely, causing incomplete decoding.
59
- try:
60
- while True:
61
- time.sleep(1)
62
- except KeyboardInterrupt:
63
- print("Script stopped by user.")
64
-
65
-
66
-def run_decode(prefill_done):
67
- # We use GPU 1 for decode node.
68
- os.environ["CUDA_VISIBLE_DEVICES"] = "1"
69
-
70
- prompts = [
71
- "Hello, my name is",
72
- "Hi, your name is",
73
- "Tell me a very long story",
74
- ]
75
- sampling_params = SamplingParams(temperature=0, top_p=0.95)
76
-
77
- # Using P2pNcclConnector to transmit KV caches between vLLM instances.
78
- # This instance is the decode node (kv_consumer, rank 1).
79
- # The number of parallel instances for KV cache transfer is set to 2,
80
- # as required for P2pNcclConnector.
81
- ktc = KVTransferConfig(
82
- kv_connector="P2pNcclConnector",
83
- kv_role="kv_consumer",
84
- kv_rank=1,
85
- kv_parallel_size=2,
86
- )
87
-
88
- # Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
89
- # memory. You may need to adjust the value to fit your GPU.
90
- llm = LLM(
91
- model="meta-llama/Meta-Llama-3.1-8B-Instruct",
92
- kv_transfer_config=ktc,
93
- max_model_len=2000,
94
- gpu_memory_utilization=0.8,
95
- )
96
-
97
- # Wait for the producer to start the pipe
98
- print("Waiting for prefill node to finish...")
99
- prefill_done.wait()
100
-
101
- # At this point when the prefill_done is set, the kv-cache should have been
102
- # transferred to this decode node, so we can start decoding.
103
- outputs = llm.generate(prompts, sampling_params)
104
- for output in outputs:
105
- prompt = output.prompt
106
- generated_text = output.outputs[0].text
107
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
108
-
109
-
110
-def main():
111
- prefill_done = Event()
112
- prefill_process = Process(target=run_prefill, args=(prefill_done,))
113
- decode_process = Process(target=run_decode, args=(prefill_done,))
114
-
115
- # Start prefill node
116
- prefill_process.start()
117
-
118
- # Start decode node
119
- decode_process.start()
120
-
121
- # Terminate the prefill node when decode is finished
122
- decode_process.join()
123
- prefill_process.terminate()
124
-
125
-
126
-if __name__ == "__main__":
127
- main()
examples/disaggregated/disaggregated_prefill.sh
deleted
-125
@@ -1,125 +0,0 @@
1
-#!/bin/bash
2
-# This file demonstrates the example usage of disaggregated prefilling
3
-# We will launch 2 vllm instances (1 for prefill and 1 for decode),
4
-# and then transfer the KV cache between them.
5
-
6
-set -xe
7
-
8
-echo "🚧🚧 Warning: The usage of disaggregated prefill is experimental and subject to change 🚧🚧"
9
-sleep 1
10
-
11
-# meta-llama/Meta-Llama-3.1-8B-Instruct or deepseek-ai/DeepSeek-V2-Lite
12
-MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
13
-
14
-# Trap the SIGINT signal (triggered by Ctrl+C)
15
-trap 'cleanup' INT
16
-
17
-# Cleanup function
18
-cleanup() {
19
- echo "Caught Ctrl+C, cleaning up..."
20
- # Cleanup commands
21
- pgrep python | xargs kill -9
22
- pkill -f python
23
- echo "Cleanup complete. Exiting."
24
- exit 0
25
-}
26
-
27
-
28
-if [[ -z "${VLLM_HOST_IP:-}" ]]; then
29
- export VLLM_HOST_IP=127.0.0.1
30
- echo "Using default VLLM_HOST_IP=127.0.0.1 (override by exporting VLLM_HOST_IP before running this script)"
31
-else
32
- echo "Using provided VLLM_HOST_IP=${VLLM_HOST_IP}"
33
-fi
34
-
35
-
36
-# install quart first -- required for disagg prefill proxy serve
37
-if python3 -c "import quart" &> /dev/null; then
38
- echo "Quart is already installed."
39
-else
40
- echo "Quart is not installed. Installing..."
41
- python3 -m pip install quart
42
-fi
43
-
44
-# a function that waits vLLM server to start
45
-wait_for_server() {
46
- local port=$1
47
- timeout 1200 bash -c "
48
- until curl -i localhost:${port}/v1/models > /dev/null; do
49
- sleep 1
50
- done" && return 0 || return 1
51
-}
52
-
53
-
54
-# You can also adjust --kv-ip and --kv-port for distributed inference.
55
-
56
-# prefilling instance, which is the KV producer
57
-CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
58
- --host 0.0.0.0 \
59
- --port 8100 \
60
- --max-model-len 100 \
61
- --gpu-memory-utilization 0.8 \
62
- --trust-remote-code \
63
- --kv-transfer-config \
64
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
65
-
66
-# decoding instance, which is the KV consumer
67
-CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
68
- --host 0.0.0.0 \
69
- --port 8200 \
70
- --max-model-len 100 \
71
- --gpu-memory-utilization 0.8 \
72
- --trust-remote-code \
73
- --kv-transfer-config \
74
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":"1e10","kv_port":"14580","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8200","send_type":"PUT_ASYNC"}}' &
75
-
76
-# wait until prefill and decode instances are ready
77
-wait_for_server 8100
78
-wait_for_server 8200
79
-
80
-# launch a proxy server that opens the service at port 8000
81
-# the workflow of this proxy:
82
-# - send the request to prefill vLLM instance (port 8100), change max_tokens
83
-# to 1
84
-# - after the prefill vLLM finishes prefill, send the request to decode vLLM
85
-# instance
86
-# NOTE: the usage of this API is subject to change --- in the future we will
87
-# introduce "vllm connect" to connect between prefill and decode instances
88
-python3 ../../benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py &
89
-sleep 1
90
-
91
-# serve two example requests
92
-output1=$(curl -X POST -s http://localhost:8000/v1/completions \
93
--H "Content-Type: application/json" \
94
--d '{
95
-"model": "'"$MODEL_NAME"'",
96
-"prompt": "San Francisco is a",
97
-"max_tokens": 10,
98
-"temperature": 0
99
-}')
100
-
101
-output2=$(curl -X POST -s http://localhost:8000/v1/completions \
102
--H "Content-Type: application/json" \
103
--d '{
104
-"model": "'"$MODEL_NAME"'",
105
-"prompt": "Santa Clara is a",
106
-"max_tokens": 10,
107
-"temperature": 0
108
-}')
109
-
110
-
111
-# Cleanup commands
112
-pgrep python | xargs kill -9
113
-pkill -f python
114
-
115
-echo ""
116
-
117
-sleep 1
118
-
119
-# Print the outputs of the curl requests
120
-echo ""
121
-echo "Output of first request: $output1"
122
-echo "Output of second request: $output2"
123
-
124
-echo "🎉🎉 Successfully finished 2 test requests! 🎉🎉"
125
-echo ""
examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh
deleted
-245
@@ -1,245 +0,0 @@
1
-#!/bin/bash
2
-
3
-# =============================================================================
4
-# vLLM Disaggregated Serving Script - P2P NCCL XpYd Architecture
5
-# =============================================================================
6
-# This script demonstrates disaggregated prefill and decode serving using
7
-# P2P NCCL communication. The architecture supports various XpYd configurations:
8
-#
9
-# - 1P3D: 1 Prefill server + 3 Decode servers (current default)
10
-# - 3P1D: 3 Prefill servers + 1 Decode server
11
-# - etc.
12
-#
13
-# Configuration can be customized via environment variables:
14
-# MODEL: Model to serve
15
-# PREFILL_GPUS: Comma-separated GPU IDs for prefill servers
16
-# DECODE_GPUS: Comma-separated GPU IDs for decode servers
17
-# PREFILL_PORTS: Comma-separated ports for prefill servers
18
-# DECODE_PORTS: Comma-separated ports for decode servers
19
-# PROXY_PORT: Proxy server port used to setup XpYd connection.
20
-# TIMEOUT_SECONDS: Server startup timeout
21
-# =============================================================================
22
-
23
-# Configuration - can be overridden via environment variables
24
-MODEL=${MODEL:-meta-llama/Llama-3.1-8B-Instruct}
25
-TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200}
26
-PROXY_PORT=${PROXY_PORT:-30001}
27
-
28
-# Default 1P3D configuration (1 Prefill + 3 Decode)
29
-PREFILL_GPUS=${PREFILL_GPUS:-0}
30
-DECODE_GPUS=${DECODE_GPUS:-1,2,3}
31
-PREFILL_PORTS=${PREFILL_PORTS:-20003}
32
-DECODE_PORTS=${DECODE_PORTS:-20005,20007,20009}
33
-
34
-echo "Warning: P2P NCCL disaggregated prefill XpYd support for vLLM v1 is experimental and subject to change."
35
-echo ""
36
-echo "Architecture Configuration:"
37
-echo " Model: $MODEL"
38
-echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS"
39
-echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS"
40
-echo " Proxy Port: $PROXY_PORT"
41
-echo " Timeout: ${TIMEOUT_SECONDS}s"
42
-echo ""
43
-
44
-PIDS=()
45
-
46
-# Switch to the directory of the current script
47
-cd "$(dirname "${BASH_SOURCE[0]}")"
48
-
49
-check_required_files() {
50
- local files=("disagg_proxy_p2p_nccl_xpyd.py")
51
- for file in "${files[@]}"; do
52
- if [[ ! -f "$file" ]]; then
53
- echo "Required file $file not found in $(pwd)"
54
- exit 1
55
- fi
56
- done
57
-}
58
-
59
-check_hf_token() {
60
- if [ -z "$HF_TOKEN" ]; then
61
- echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
62
- echo "Example: export HF_TOKEN=your_token_here"
63
- exit 1
64
- fi
65
- if [[ "$HF_TOKEN" != hf_* ]]; then
66
- echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
67
- exit 1
68
- fi
69
- echo "HF_TOKEN is set and valid."
70
-}
71
-
72
-check_num_gpus() {
73
- # Check if the number of GPUs are >=2 via nvidia-smi
74
- num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
75
- if [ "$num_gpus" -lt 2 ]; then
76
- echo "You need at least 2 GPUs to run disaggregated prefill."
77
- exit 1
78
- else
79
- echo "Found $num_gpus GPUs."
80
- fi
81
-}
82
-
83
-ensure_python_library_installed() {
84
- echo "Checking if $1 is installed..."
85
- if ! python3 -c "import $1" > /dev/null 2>&1; then
86
- echo "$1 is not installed. Please install it via pip install $1."
87
- exit 1
88
- else
89
- echo "$1 is installed."
90
- fi
91
-}
92
-
93
-cleanup() {
94
- echo "Stopping everything…"
95
- trap - INT TERM # prevent re-entrancy
96
- pkill -9 -f "disagg_proxy_p2p_nccl_xpyd.py"
97
- kill -- -$$ # negative PID == "this whole process-group"
98
- wait # reap children so we don't leave zombies
99
- exit 0
100
-}
101
-
102
-wait_for_server() {
103
- local port=$1
104
- local timeout_seconds=$TIMEOUT_SECONDS
105
- local start_time=$(date +%s)
106
-
107
- echo "Waiting for server on port $port..."
108
-
109
- while true; do
110
- if curl -s "localhost:${port}/v1/completions" > /dev/null; then
111
- echo "Server on port $port is ready."
112
- return 0
113
- fi
114
-
115
- local now=$(date +%s)
116
- if (( now - start_time >= timeout_seconds )); then
117
- echo "Timeout waiting for server on port $port"
118
- return 1
119
- fi
120
-
121
- sleep 1
122
- done
123
-}
124
-
125
-main() {
126
- check_required_files
127
- check_hf_token
128
- check_num_gpus
129
- ensure_python_library_installed pandas
130
- ensure_python_library_installed datasets
131
- ensure_python_library_installed vllm
132
- ensure_python_library_installed quart
133
-
134
- trap cleanup INT
135
- trap cleanup USR1
136
- trap cleanup TERM
137
-
138
- echo "Launching disaggregated serving components..."
139
- echo "Please check the log files for detailed output:"
140
- echo " - prefill*.log: Prefill server logs"
141
- echo " - decode*.log: Decode server logs"
142
- echo " - proxy.log: Proxy server log"
143
-
144
- # =============================================================================
145
- # Launch Proxy Server
146
- # =============================================================================
147
- echo ""
148
- echo "Starting proxy server on port $PROXY_PORT..."
149
- python3 disagg_proxy_p2p_nccl_xpyd.py &
150
- PIDS+=($!)
151
-
152
- # Parse GPU and port arrays
153
- IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS"
154
- IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS"
155
- IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS"
156
- IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
157
-
158
- # =============================================================================
159
- # Launch Prefill Servers (X Producers)
160
- # =============================================================================
161
- echo ""
162
- echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..."
163
- for i in "${!PREFILL_GPU_ARRAY[@]}"; do
164
- local gpu_id=${PREFILL_GPU_ARRAY[$i]}
165
- local port=${PREFILL_PORT_ARRAY[$i]}
166
- local kv_port=$((21001 + i))
167
-
168
- echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
169
- CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
170
- --enforce-eager \
171
- --host 0.0.0.0 \
172
- --port "$port" \
173
- --tensor-parallel-size 1 \
174
- --seed 1024 \
175
- --dtype float16 \
176
- --max-model-len 10000 \
177
- --max-num-batched-tokens 10000 \
178
- --max-num-seqs 256 \
179
- --trust-remote-code \
180
- --gpu-memory-utilization 0.9 \
181
- --kv-transfer-config \
182
- "{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"1e1\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > prefill$((i+1)).log 2>&1 &
183
- PIDS+=($!)
184
- done
185
-
186
- # =============================================================================
187
- # Launch Decode Servers (Y Decoders)
188
- # =============================================================================
189
- echo ""
190
- echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..."
191
- for i in "${!DECODE_GPU_ARRAY[@]}"; do
192
- local gpu_id=${DECODE_GPU_ARRAY[$i]}
193
- local port=${DECODE_PORT_ARRAY[$i]}
194
- local kv_port=$((22001 + i))
195
-
196
- echo " Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
197
- CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
198
- --enforce-eager \
199
- --host 0.0.0.0 \
200
- --port "$port" \
201
- --tensor-parallel-size 1 \
202
- --seed 1024 \
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.7 \
209
- --kv-transfer-config \
210
- "{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > decode$((i+1)).log 2>&1 &
211
- PIDS+=($!)
212
- done
213
-
214
- # =============================================================================
215
- # Wait for All Servers to Start
216
- # =============================================================================
217
- echo ""
218
- echo "Waiting for all servers to start..."
219
- for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
220
- if ! wait_for_server "$port"; then
221
- echo "Failed to start server on port $port"
222
- cleanup
223
- # shellcheck disable=SC2317
224
- exit 1
225
- fi
226
- done
227
-
228
- echo ""
229
- echo "All servers are up. Starting benchmark..."
230
-
231
- # =============================================================================
232
- # Run Benchmark
233
- # =============================================================================
234
- cd ../../../benchmarks/
235
- vllm bench serve --port 10001 --seed "$(date +%s)" \
236
- --model "$MODEL" \
237
- --dataset-name random --random-input-len 7500 --random-output-len 200 \
238
- --num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
239
-
240
- echo "Benchmarking done. Cleaning up..."
241
-
242
- cleanup
243
-}
244
-
245
-main
examples/disaggregated/p2p_nccl_xpyd/disagg_proxy_p2p_nccl_xpyd.py
deleted
-190
@@ -1,190 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
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
13
-import zmq
14
-from quart import Quart, make_response, request
15
-
16
-count = 0
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:
39
- socks = dict(poller.poll())
40
- if router_socket in socks:
41
- remote_address, message = router_socket.recv_multipart()
42
- # data: {"type": "P", "http_address": "ip:port",
43
- # "zmq_address": "ip:port"}
44
- data = msgpack.loads(message)
45
- if data["type"] == "P":
46
- global prefill_instances
47
- global prefill_cv
48
- with prefill_cv:
49
- node = prefill_instances.get(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:
60
- node = decode_instances.get(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",
69
- remote_address,
70
- data,
71
- )
72
- return
73
-
74
- if node is None:
75
- print(f"🔵Add [HTTP:{data['http_address']}, ZMQ:{data['zmq_address']}]")
76
-
77
-
78
-def start_service_discovery(hostname, port):
79
- if not hostname:
80
- hostname = socket.gethostname()
81
- if port == 0:
82
- raise ValueError("Port cannot be 0")
83
-
84
- context = zmq.Context()
85
- router_socket = context.socket(zmq.ROUTER)
86
- router_socket.bind(f"tcp://{hostname}:{port}")
87
-
88
- poller = zmq.Poller()
89
- poller.register(router_socket, zmq.POLLIN)
90
-
91
- _listener_thread = threading.Thread(
92
- target=_listen_for_register, args=[poller, router_socket], daemon=True
93
- )
94
- _listener_thread.start()
95
- return _listener_thread
96
-
97
-
98
-AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
99
-
100
-app = Quart(__name__)
101
-
102
-
103
-def random_uuid() -> str:
104
- return str(uuid.uuid4().hex)
105
-
106
-
107
-async def forward_request(url, data, request_id):
108
- async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
109
- headers = {
110
- "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
111
- "X-Request-Id": request_id,
112
- }
113
- async with session.post(url=url, json=data, headers=headers) as response:
114
- if response.status == 200:
115
- if True:
116
- async for chunk_bytes in response.content.iter_chunked(1024):
117
- yield chunk_bytes
118
- else:
119
- content = await response.read()
120
- yield content
121
-
122
-
123
-@app.route("/v1/completions", methods=["POST"])
124
-@app.route("/v1/chat/completions", methods=["POST"])
125
-async def handle_request():
126
- try:
127
- original_request_data = await request.get_json()
128
-
129
- prefill_request = original_request_data.copy()
130
- # change max_tokens = 1 to let it only do prefill
131
- prefill_request["max_tokens"] = 1
132
- if "max_completion_tokens" in prefill_request:
133
- prefill_request["max_completion_tokens"] = 1
134
-
135
- global count
136
- global prefill_instances
137
- global prefill_cv
138
- with prefill_cv:
139
- prefill_list = list(prefill_instances.items())
140
- prefill_addr, prefill_zmq_addr = prefill_list[count % len(prefill_list)]
141
- prefill_zmq_addr = prefill_zmq_addr[0]
142
-
143
- global decode_instances
144
- global decode_cv
145
- with decode_cv:
146
- decode_list = list(decode_instances.items())
147
- decode_addr, decode_zmq_addr = decode_list[count % len(decode_list)]
148
- decode_zmq_addr = decode_zmq_addr[0]
149
-
150
- print(
151
- f"handle_request count: {count}, [HTTP:{prefill_addr}, "
152
- f"ZMQ:{prefill_zmq_addr}] 👉 [HTTP:{decode_addr}, "
153
- f"ZMQ:{decode_zmq_addr}]"
154
- )
155
- count += 1
156
-
157
- request_id = (
158
- f"___prefill_addr_{prefill_zmq_addr}___decode_addr_"
159
- f"{decode_zmq_addr}_{random_uuid()}"
160
- )
161
-
162
- # finish prefill
163
- async for _ in forward_request(
164
- f"http://{prefill_addr}{request.path}", prefill_request, request_id
165
- ):
166
- continue
167
-
168
- # return decode
169
- generator = forward_request(
170
- f"http://{decode_addr}{request.path}", original_request_data, request_id
171
- )
172
- response = await make_response(generator)
173
- response.timeout = None
174
-
175
- return response
176
-
177
- except Exception as e:
178
- import sys
179
- import traceback
180
-
181
- exc_info = sys.exc_info()
182
- print("Error occurred in disagg prefill proxy server")
183
- print(e)
184
- print("".join(traceback.format_exception(*exc_info)))
185
-
186
-
187
-if __name__ == "__main__":
188
- t = start_service_discovery("0.0.0.0", 30001)
189
- app.run(host="0.0.0.0", port=10001)
190
- t.join()
vllm/config/kv_transfer.py
+1
-2
@@ -48,8 +48,7 @@ class KVTransferConfig:
48
Currently only 1P1D is supported."""
49
50
kv_parallel_size: int = 1
51
- """The number of parallel instances for KV cache transfer. For
52
- P2pNcclConnector, this should be 2."""
51
+ """The number of parallel instances for KV cache transfer."""
52
53
kv_ip: str = "127.0.0.1"
54
"""The KV connector ip, used to build distributed connection."""
vllm/distributed/kv_transfer/kv_connector/factory.py
-6
@@ -161,12 +161,6 @@ KVConnectorFactory.register_connector(
161
"ExampleHiddenStatesConnector",
162
)
163
164
-KVConnectorFactory.register_connector(
165
- "P2pNcclConnector",
166
- "vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_connector",
167
- "P2pNcclConnector",
168
-)
169
-
164
KVConnectorFactory.register_connector(
165
"LMCacheConnectorV1",
166
"vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+1
-1
@@ -324,7 +324,7 @@ class MoRIIOConstants:
324
DEFAULT_DEFER_TIMEOUT = 60.0
325
326
327
-# The router embeds both zmq_addresses in the request_id (similar to P2pNcclConnector):
327
+# The router embeds both zmq_addresses in the request_id:
328
# "___prefill_addr_{zmq}___decode_addr_{zmq}_{32-hex-uuid}"
329
# MoRIIO zmq_address format: "host:IP,handshake:PORT,notify:PORT"
330
#
vllm/distributed/kv_transfer/kv_connector/v1/p2p/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
deleted
-478
@@ -1,478 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, Any
6
-
7
-import regex as re
8
-import torch
9
-
10
-from vllm.config import VllmConfig
11
-from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12
- KVConnectorBase_V1,
13
- KVConnectorMetadata,
14
- KVConnectorRole,
15
-)
16
-from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
17
- P2pNcclEngine,
18
-)
19
-from vllm.distributed.parallel_state import get_world_group
20
-from vllm.logger import init_logger
21
-from vllm.v1.attention.backend import AttentionMetadata
22
-from vllm.v1.core.sched.output import SchedulerOutput
23
-
24
-if TYPE_CHECKING:
25
- from vllm.forward_context import ForwardContext
26
- from vllm.v1.core.kv_cache_manager import KVCacheBlocks
27
- from vllm.v1.kv_cache_interface import KVCacheConfig
28
- from vllm.v1.request import Request
29
-
30
-logger = init_logger(__name__)
31
-
32
-
33
-@dataclass
34
-class ReqMeta:
35
- # Request Id
36
- request_id: str
37
- # Request block ids
38
- block_ids: torch.Tensor
39
- # Request num tokens
40
- num_tokens: int
41
-
42
- @staticmethod
43
- def make_meta(
44
- request_id: str, token_ids: list[int], block_ids: list[int], block_size: int
45
- ) -> "ReqMeta":
46
- block_ids_tensor = torch.tensor(block_ids)
47
- return ReqMeta(
48
- request_id=request_id,
49
- block_ids=block_ids_tensor,
50
- num_tokens=len(token_ids),
51
- )
52
-
53
-
54
-@dataclass
55
-class P2pNcclConnectorMetadata(KVConnectorMetadata):
56
- requests: list[ReqMeta]
57
-
58
- def __init__(self):
59
- self.requests = []
60
-
61
- def add_request(
62
- self,
63
- request_id: str,
64
- token_ids: list[int],
65
- block_ids: list[int],
66
- block_size: int,
67
- ) -> None:
68
- self.requests.append(
69
- ReqMeta.make_meta(request_id, token_ids, block_ids, block_size)
70
- )
71
-
72
-
73
-class P2pNcclConnector(KVConnectorBase_V1):
74
- def __init__(
75
- self,
76
- vllm_config: "VllmConfig",
77
- role: KVConnectorRole,
78
- kv_cache_config: "KVCacheConfig",
79
- ):
80
- super().__init__(
81
- vllm_config=vllm_config,
82
- role=role,
83
- kv_cache_config=kv_cache_config,
84
- )
85
- self._block_size = vllm_config.cache_config.block_size
86
- self._requests_need_load: dict[str, Any] = {}
87
- self.is_producer = self._kv_transfer_config.is_kv_producer
88
- self.chunked_prefill: dict[str, tuple[list[int], list[int] | None]] = {}
89
-
90
- self._rank = get_world_group().rank if role == KVConnectorRole.WORKER else 0
91
- self._local_rank = (
92
- get_world_group().local_rank if role == KVConnectorRole.WORKER else 0
93
- )
94
-
95
- self.p2p_nccl_engine = (
96
- P2pNcclEngine(
97
- local_rank=self._local_rank,
98
- config=self._kv_transfer_config,
99
- hostname="",
100
- port_offset=self._rank,
101
- )
102
- if role == KVConnectorRole.WORKER
103
- else None
104
- )
105
-
106
- # ==============================
107
- # Worker-side methods
108
- # ==============================
109
-
110
- def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
111
- """Start loading the KV cache from the connector buffer to vLLM's
112
- paged KV buffer.
113
-
114
- Args:
115
- forward_context (ForwardContext): the forward context.
116
- **kwargs: additional arguments for the load operation
117
-
118
- Note:
119
- The number of elements in kv_caches and layer_names should be
120
- the same.
121
- """
122
-
123
- # Only consumer/decode loads KV Cache
124
- if self.is_producer:
125
- return
126
-
127
- assert self.p2p_nccl_engine is not None
128
-
129
- attn_metadata = forward_context.attn_metadata
130
- if attn_metadata is None:
131
- return
132
-
133
- def inject_kv_into_layer(
134
- layer: torch.Tensor,
135
- kv_cache: torch.Tensor,
136
- block_ids: torch.Tensor,
137
- request_id: str,
138
- ) -> None:
139
- """
140
- Inject KV cache data into a given attention layer tensor.
141
-
142
- This function updates `layer` in-place with values from `kv_cache`.
143
- All backends (MLA, FlashAttention, FlashInfer, TritonAttention)
144
- are indexed along the first dimension (block index).
145
-
146
- If the number of provided block IDs does not match the number of KV
147
- blocks, only the overlapping portion is updated, and a warning is
148
- logged.
149
-
150
- Args:
151
- layer (torch.Tensor): The attention layer KV tensor to update.
152
- kv_cache (torch.Tensor): The KV cache tensor to inject.
153
- block_ids (torch.Tensor): Indices of the blocks to update.
154
- request_id (str): Request identifier used for logging.
155
-
156
- Returns:
157
- None. The function modifies `layer` in-place.
158
- """
159
- num_block = kv_cache.shape[0]
160
- self.check_tensors_except_dim(layer, kv_cache, 0)
161
- if len(block_ids) == num_block:
162
- layer[block_ids, ...] = kv_cache
163
- else:
164
- layer[block_ids[:num_block], ...] = kv_cache
165
- logger.warning(
166
- "🚧kv_cache does not match, block_ids:%d, "
167
- "num_block:%d, request_id:%s",
168
- len(block_ids),
169
- num_block,
170
- request_id,
171
- )
172
-
173
- # Get the metadata
174
- metadata: KVConnectorMetadata = self._get_connector_metadata()
175
- assert isinstance(metadata, P2pNcclConnectorMetadata)
176
-
177
- if metadata is None:
178
- return
179
-
180
- # Load the KV for each request each layer
181
- for request in metadata.requests:
182
- request_id = request.request_id
183
- ip, port = self.parse_request_id(request_id, False)
184
- remote_address = ip + ":" + str(port + self._rank)
185
- for layer_name in forward_context.no_compile_layers:
186
- layer = forward_context.no_compile_layers[layer_name]
187
-
188
- # Only process layers that have kv_cache
189
- # attribute (attention layers) Skip non-attention
190
- # layers like FusedMoE
191
- kv_cache = getattr(layer, "kv_cache", None)
192
- if kv_cache is None:
193
- continue
194
-
195
- layer = kv_cache
196
-
197
- kv_cache = self.p2p_nccl_engine.recv_tensor(
198
- request.request_id + "#" + layer_name, remote_address
199
- )
200
-
201
- if kv_cache is None:
202
- logger.warning("🚧kv_cache is None, %s", request.request_id)
203
- continue
204
-
205
- inject_kv_into_layer(
206
- layer, kv_cache, request.block_ids, request.request_id
207
- )
208
-
209
- def wait_for_layer_load(self, layer_name: str) -> None:
210
- """Blocking until the KV for a specific layer is loaded into vLLM's
211
- paged buffer.
212
-
213
- This interface will be useful for layer-by-layer pipelining.
214
-
215
- Args:
216
- layer_name: the name of that layer
217
- """
218
- return
219
-
220
- def save_kv_layer(
221
- self,
222
- layer_name: str,
223
- kv_layer: torch.Tensor,
224
- attn_metadata: AttentionMetadata,
225
- **kwargs: Any,
226
- ) -> None:
227
- """Start saving the KV cache of the layer from vLLM's paged buffer
228
- to the connector.
229
-
230
- Args:
231
- layer_name (str): the name of the layer.
232
- kv_layer (torch.Tensor): the paged KV buffer of the current
233
- layer in vLLM.
234
- attn_metadata (AttentionMetadata): the attention metadata.
235
- **kwargs: additional arguments for the save operation.
236
- """
237
-
238
- # Only producer/prefill saves KV Cache
239
- if not self.is_producer:
240
- return
241
-
242
- assert self.p2p_nccl_engine is not None
243
-
244
- connector_metadata = self._get_connector_metadata()
245
- assert isinstance(connector_metadata, P2pNcclConnectorMetadata)
246
- for request in connector_metadata.requests:
247
- request_id = request.request_id
248
- ip, port = self.parse_request_id(request_id, True)
249
- remote_address = ip + ":" + str(port + self._rank)
250
-
251
- kv_cache = kv_layer[request.block_ids, ...]
252
- self.p2p_nccl_engine.send_tensor(
253
- request_id + "#" + layer_name, kv_cache, remote_address
254
- )
255
-
256
- def wait_for_save(self):
257
- if self.is_producer:
258
- assert self.p2p_nccl_engine is not None
259
- self.p2p_nccl_engine.wait_for_sent()
260
-
261
- def get_finished(
262
- self, finished_req_ids: set[str], **kwargs: Any
263
- ) -> tuple[set[str] | None, set[str] | None]:
264
- """
265
- Notifies worker-side connector ids of requests that have
266
- finished generating tokens.
267
-
268
- Returns:
269
- ids of requests that have finished asynchronous transfer,
270
- tuple of (sending/saving ids, recving/loading ids).
271
- The finished saves/sends req ids must belong to a set provided in a
272
- call to this method (this call or a prior one).
273
- """
274
-
275
- assert self.p2p_nccl_engine is not None
276
-
277
- no_compile_layers = self._vllm_config.compilation_config.static_forward_context
278
- return self.p2p_nccl_engine.get_finished(finished_req_ids, no_compile_layers)
279
-
280
- # ==============================
281
- # Scheduler-side methods
282
- # ==============================
283
-
284
- def get_num_new_matched_tokens(
285
- self,
286
- request: "Request",
287
- num_computed_tokens: int,
288
- ) -> tuple[int, bool]:
289
- """
290
- Get number of new tokens that can be loaded from the
291
- external KV cache beyond the num_computed_tokens.
292
-
293
- Args:
294
- request (Request): the request object.
295
- num_computed_tokens (int): the number of locally
296
- computed tokens for this request
297
-
298
- Returns:
299
- the number of tokens that can be loaded from the
300
- external KV cache beyond what is already computed.
301
- """
302
- if self.is_producer:
303
- return 0, False
304
-
305
- prompt_token_ids = request.prompt_token_ids or []
306
- num_external_tokens = len(prompt_token_ids) - 1 - num_computed_tokens
307
-
308
- if num_external_tokens < 0:
309
- num_external_tokens = 0
310
-
311
- return num_external_tokens, False
312
-
313
- def update_state_after_alloc(
314
- self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
315
- ):
316
- """
317
- Update KVConnector state after block allocation.
318
- """
319
- if not self.is_producer and num_external_tokens > 0:
320
- self._requests_need_load[request.request_id] = (
321
- request,
322
- blocks.get_block_ids()[0],
323
- )
324
-
325
- def build_connector_meta(
326
- self,
327
- scheduler_output: SchedulerOutput,
328
- ) -> KVConnectorMetadata:
329
- """Build the connector metadata for this step.
330
-
331
- This function should NOT modify any fields in the scheduler_output.
332
- Also, calling this function will reset the state of the connector.
333
-
334
- Args:
335
- scheduler_output (SchedulerOutput): the scheduler output object.
336
- """
337
-
338
- meta = P2pNcclConnectorMetadata()
339
-
340
- for new_req in scheduler_output.scheduled_new_reqs:
341
- if self.is_producer:
342
- num_scheduled_tokens = (scheduler_output.num_scheduled_tokens)[
343
- new_req.req_id
344
- ]
345
- num_tokens = num_scheduled_tokens + new_req.num_computed_tokens
346
- # the request's prompt is chunked prefill
347
- if num_tokens < len(new_req.prompt_token_ids or []):
348
- # 'CachedRequestData' has no attribute 'prompt_token_ids'
349
- self.chunked_prefill[new_req.req_id] = (
350
- new_req.block_ids[0],
351
- new_req.prompt_token_ids,
352
- )
353
- continue
354
- # the request's prompt is not chunked prefill
355
- meta.add_request(
356
- request_id=new_req.req_id,
357
- token_ids=new_req.prompt_token_ids or [],
358
- block_ids=new_req.block_ids[0],
359
- block_size=self._block_size,
360
- )
361
- continue
362
- if new_req.req_id in self._requests_need_load:
363
- meta.add_request(
364
- request_id=new_req.req_id,
365
- token_ids=new_req.prompt_token_ids or [],
366
- block_ids=new_req.block_ids[0],
367
- block_size=self._block_size,
368
- )
369
- self._requests_need_load.pop(new_req.req_id)
370
-
371
- cached_reqs = scheduler_output.scheduled_cached_reqs
372
- for i, req_id in enumerate(cached_reqs.req_ids):
373
- num_computed_tokens = cached_reqs.num_computed_tokens[i]
374
- new_block_ids = cached_reqs.new_block_ids[i]
375
- resumed_from_preemption = req_id in cached_reqs.resumed_req_ids
376
-
377
- if self.is_producer:
378
- num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
379
- num_tokens = num_scheduled_tokens + num_computed_tokens
380
- assert req_id in self.chunked_prefill
381
- assert new_block_ids is not None
382
- block_ids = new_block_ids[0]
383
- if not resumed_from_preemption:
384
- block_ids = self.chunked_prefill[req_id][0] + block_ids
385
- prompt_token_ids = self.chunked_prefill[req_id][1]
386
- assert prompt_token_ids is not None
387
- # the request's prompt is chunked prefill again
388
- if num_tokens < len(prompt_token_ids):
389
- self.chunked_prefill[req_id] = (block_ids, prompt_token_ids)
390
- continue
391
- # the request's prompt is all prefilled finally
392
- meta.add_request(
393
- request_id=req_id,
394
- token_ids=prompt_token_ids,
395
- block_ids=block_ids,
396
- block_size=self._block_size,
397
- )
398
- self.chunked_prefill.pop(req_id, None)
399
- continue
400
-
401
- # NOTE(rob): here we rely on the resumed requests being
402
- # the first N requests in the list scheduled_cache_reqs.
403
- if not resumed_from_preemption:
404
- break
405
- if req_id in self._requests_need_load:
406
- request, _ = self._requests_need_load.pop(req_id)
407
- total_tokens = num_computed_tokens + 1
408
- token_ids = request.all_token_ids[:total_tokens]
409
-
410
- # NOTE(rob): For resumed req, new_block_ids is all
411
- # of the block_ids for the request.
412
- assert new_block_ids is not None
413
- block_ids = new_block_ids[0]
414
-
415
- meta.add_request(
416
- request_id=req_id,
417
- token_ids=token_ids,
418
- block_ids=block_ids,
419
- block_size=self._block_size,
420
- )
421
-
422
- self._requests_need_load.clear()
423
- return meta
424
-
425
- def request_finished(
426
- self,
427
- request: "Request",
428
- block_ids: list[int],
429
- ) -> tuple[bool, dict[str, Any] | None]:
430
- """
431
- Called when a request has finished, before its blocks are freed.
432
-
433
- Returns:
434
- True if the request is being saved/sent asynchronously and blocks
435
- should not be freed until the request_id is returned from
436
- get_finished().
437
- Optional KVTransferParams to be included in the request outputs
438
- returned by the engine.
439
- """
440
-
441
- self.chunked_prefill.pop(request.request_id, None)
442
-
443
- return False, None
444
-
445
- # ==============================
446
- # Static methods
447
- # ==============================
448
-
449
- @staticmethod
450
- def parse_request_id(request_id: str, is_prefill=True) -> tuple[str, int]:
451
- # Regular expression to match the string hostname and integer port
452
- if is_prefill:
453
- pattern = r"___decode_addr_(.*):(\d+)"
454
- else:
455
- pattern = r"___prefill_addr_(.*):(\d+)___"
456
-
457
- # Use re.search to find the pattern in the request_id
458
- match = re.search(pattern, request_id)
459
- if match:
460
- # Extract the ranks
461
- ip = match.group(1)
462
- port = int(match.group(2))
463
-
464
- return ip, port
465
- raise ValueError(f"Request id {request_id} does not contain hostname and port")
466
-
467
- @staticmethod
468
- def check_tensors_except_dim(tensor1, tensor2, dim):
469
- shape1 = tensor1.size()
470
- shape2 = tensor2.size()
471
-
472
- if len(shape1) != len(shape2) or not all(
473
- s1 == s2 for i, (s1, s2) in enumerate(zip(shape1, shape2)) if i != dim
474
- ):
475
- raise NotImplementedError(
476
- "Currently, only symmetric TP is supported. Asymmetric TP, PP,"
477
- "and others will be supported in future PRs."
478
- )
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
deleted
-632
@@ -1,632 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import json
5
-import logging
6
-import os
7
-import threading
8
-import time
9
-from collections import deque
10
-from contextlib import contextmanager
11
-from dataclasses import dataclass
12
-from typing import Any
13
-
14
-import msgpack
15
-import torch
16
-import zmq
17
-
18
-from vllm.config.kv_transfer import KVTransferConfig
19
-from vllm.distributed.device_communicators.pynccl_wrapper import (
20
- NCCLLibrary,
21
- buffer_type,
22
- cudaStream_t,
23
- ncclComm_t,
24
- ncclDataTypeEnum,
25
-)
26
-from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import ( # noqa: E501
27
- TensorMemoryPool,
28
-)
29
-from vllm.utils.network_utils import get_ip
30
-from vllm.utils.torch_utils import current_stream
31
-
32
-logger = logging.getLogger(__name__)
33
-
34
-DEFAULT_MEM_POOL_SIZE_GB = 32
35
-
36
-
37
-@contextmanager
38
-def set_p2p_nccl_context(num_channels: str):
39
- original_values: dict[str, Any] = {}
40
- env_vars = [
41
- "NCCL_MAX_NCHANNELS",
42
- "NCCL_MIN_NCHANNELS",
43
- "NCCL_CUMEM_ENABLE",
44
- "NCCL_BUFFSIZE",
45
- "NCCL_PROTO", # LL,LL128,SIMPLE
46
- "NCCL_ALGO", # RING,TREE
47
- ]
48
-
49
- for var in env_vars:
50
- original_values[var] = os.environ.get(var)
51
-
52
- logger.info("set_p2p_nccl_context, original_values: %s", original_values)
53
-
54
- try:
55
- os.environ["NCCL_MAX_NCHANNELS"] = num_channels
56
- os.environ["NCCL_MIN_NCHANNELS"] = num_channels
57
- os.environ["NCCL_CUMEM_ENABLE"] = "1"
58
- yield
59
- finally:
60
- for var in env_vars:
61
- if original_values[var] is not None:
62
- os.environ[var] = original_values[var]
63
- else:
64
- os.environ.pop(var, None)
65
-
66
-
67
-@dataclass
68
-class SendQueueItem:
69
- tensor_id: str
70
- remote_address: str
71
- tensor: torch.Tensor
72
-
73
-
74
-class P2pNcclEngine:
75
- def __init__(
76
- self,
77
- local_rank: int,
78
- config: KVTransferConfig,
79
- hostname: str = "",
80
- port_offset: int = 0,
81
- library_path: str | None = None,
82
- ) -> None:
83
- self.config = config
84
- self.rank = port_offset
85
- self.local_rank = local_rank
86
- self.device = torch.device(f"cuda:{self.local_rank}")
87
- self.nccl = NCCLLibrary(library_path)
88
-
89
- if not hostname:
90
- hostname = get_ip()
91
- port = int(self.config.kv_port) + port_offset
92
- if port == 0:
93
- raise ValueError("Port cannot be 0")
94
- self._hostname = hostname
95
- self._port = port
96
-
97
- # Each card corresponds to a ZMQ address.
98
- self.zmq_address = f"{self._hostname}:{self._port}"
99
-
100
- # If `proxy_ip` or `proxy_port` is `""`,
101
- # then the ping thread will not be enabled.
102
- proxy_ip = self.config.get_from_extra_config("proxy_ip", "")
103
- proxy_port = self.config.get_from_extra_config("proxy_port", "")
104
- if proxy_ip == "" or proxy_port == "":
105
- self.proxy_address = ""
106
- self.http_address = ""
107
- else:
108
- self.proxy_address = proxy_ip + ":" + proxy_port
109
- # the `http_port` must be consistent with the port of OpenAI.
110
- http_port = self.config.get_from_extra_config("http_port", None)
111
- if http_port is None:
112
- example_cfg = {
113
- "kv_connector": "P2pNcclConnector",
114
- "kv_connector_extra_config": {"http_port": 8000},
115
- }
116
- example = (
117
- f"--port=8000 --kv-transfer-config='{json.dumps(example_cfg)}'"
118
- )
119
- raise ValueError(
120
- "kv_connector_extra_config.http_port is required. "
121
- f"Example: {example}"
122
- )
123
- self.http_address = f"{self._hostname}:{http_port}"
124
-
125
- self.context = zmq.Context()
126
- self.router_socket = self.context.socket(zmq.ROUTER)
127
- self.router_socket.bind(f"tcp://{self.zmq_address}")
128
-
129
- self.poller = zmq.Poller()
130
- self.poller.register(self.router_socket, zmq.POLLIN)
131
-
132
- self.send_store_cv = threading.Condition()
133
- self.send_queue_cv = threading.Condition()
134
- self.recv_store_cv = threading.Condition()
135
-
136
- self.send_stream = torch.cuda.Stream()
137
- self.recv_stream = torch.cuda.Stream()
138
-
139
- mem_pool_size_gb = float(
140
- self.config.get_from_extra_config(
141
- "mem_pool_size_gb", DEFAULT_MEM_POOL_SIZE_GB
142
- )
143
- )
144
- self.pool = TensorMemoryPool(
145
- max_block_size=int(mem_pool_size_gb * 1024**3)
146
- ) # GB
147
-
148
- # The sending type includes tree mutually exclusive options:
149
- # PUT, GET, PUT_ASYNC.
150
- self.send_type = self.config.get_from_extra_config("send_type", "PUT_ASYNC")
151
- if self.send_type == "GET":
152
- # tensor_id: torch.Tensor
153
- self.send_store: dict[str, torch.Tensor] = {}
154
- else:
155
- # PUT or PUT_ASYNC
156
- # tensor_id: torch.Tensor
157
- self.send_queue: deque[SendQueueItem] = deque()
158
- if self.send_type == "PUT_ASYNC":
159
- self._send_thread = threading.Thread(
160
- target=self.send_async, daemon=True
161
- )
162
- self._send_thread.start()
163
-
164
- # tensor_id: torch.Tensor/(addr, dtype, shape)
165
- self.recv_store: dict[str, Any] = {}
166
- self.recv_request_id_to_tensor_ids: dict[str, set[str]] = {}
167
- self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
168
- self.socks: dict[str, Any] = {} # remote_address: client socket
169
- self.comms: dict[str, Any] = {} # remote_address: (ncclComm_t, rank)
170
-
171
- self.buffer_size = 0
172
- self.buffer_size_threshold = float(self.config.kv_buffer_size)
173
-
174
- self.nccl_num_channels = self.config.get_from_extra_config(
175
- "nccl_num_channels", "8"
176
- )
177
-
178
- self._listener_thread = threading.Thread(
179
- target=self.listen_for_requests, daemon=True
180
- )
181
- self._listener_thread.start()
182
-
183
- self._ping_thread = None
184
- if port_offset == 0 and self.proxy_address != "":
185
- self._ping_thread = threading.Thread(target=self.ping, daemon=True)
186
- self._ping_thread.start()
187
-
188
- logger.info(
189
- "💯P2pNcclEngine init, rank:%d, local_rank:%d, http_address:%s, "
190
- "zmq_address:%s, proxy_address:%s, send_type:%s, buffer_size_"
191
- "threshold:%.2f, nccl_num_channels:%s",
192
- self.rank,
193
- self.local_rank,
194
- self.http_address,
195
- self.zmq_address,
196
- self.proxy_address,
197
- self.send_type,
198
- self.buffer_size_threshold,
199
- self.nccl_num_channels,
200
- )
201
-
202
- def create_connect(self, remote_address: str | None = None):
203
- assert remote_address is not None
204
- if remote_address not in self.socks:
205
- sock = self.context.socket(zmq.DEALER)
206
- sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
207
- sock.connect(f"tcp://{remote_address}")
208
- self.socks[remote_address] = sock
209
- if remote_address in self.comms:
210
- logger.info(
211
- "👋comm exists, remote_address:%s, comms:%s",
212
- remote_address,
213
- self.comms,
214
- )
215
- return sock, self.comms[remote_address]
216
-
217
- unique_id = self.nccl.ncclGetUniqueId()
218
- data = {"cmd": "NEW", "unique_id": bytes(unique_id.internal)}
219
- sock.send(msgpack.dumps(data))
220
-
221
- with torch.accelerator.device_index(self.device.index):
222
- rank = 0
223
- with set_p2p_nccl_context(self.nccl_num_channels):
224
- comm: ncclComm_t = self.nccl.ncclCommInitRank(2, unique_id, rank)
225
- self.comms[remote_address] = (comm, rank)
226
- logger.info(
227
- "🤝ncclCommInitRank Success, %s👉%s, MyRank:%s",
228
- self.zmq_address,
229
- remote_address,
230
- rank,
231
- )
232
-
233
- return self.socks[remote_address], self.comms[remote_address]
234
-
235
- def send_tensor(
236
- self,
237
- tensor_id: str,
238
- tensor: torch.Tensor,
239
- remote_address: str | None = None,
240
- ) -> bool:
241
- if remote_address is None:
242
- with self.recv_store_cv:
243
- self.recv_store[tensor_id] = tensor
244
- self.recv_store_cv.notify()
245
- return True
246
-
247
- item = SendQueueItem(
248
- tensor_id=tensor_id, remote_address=remote_address, tensor=tensor
249
- )
250
-
251
- if self.send_type == "PUT":
252
- return self.send_sync(item)
253
-
254
- if self.send_type == "PUT_ASYNC":
255
- with self.send_queue_cv:
256
- self.send_queue.append(item)
257
- self.send_queue_cv.notify()
258
- return True
259
-
260
- # GET
261
- with self.send_store_cv:
262
- tensor_size = tensor.element_size() * tensor.numel()
263
- if tensor_size > self.buffer_size_threshold:
264
- logger.warning(
265
- "❗[GET]tensor_id:%s, tensor_size:%d, is greater than"
266
- "buffer size threshold :%d, skip send to %s, rank:%d",
267
- tensor_id,
268
- tensor_size,
269
- self.buffer_size_threshold,
270
- remote_address,
271
- self.rank,
272
- )
273
- return False
274
- while self.buffer_size + tensor_size > self.buffer_size_threshold:
275
- assert len(self.send_store) > 0
276
- oldest_tensor_id = next(iter(self.send_store))
277
- oldest_tensor = self.send_store.pop(oldest_tensor_id)
278
- oldest_tensor_size = (
279
- oldest_tensor.element_size() * oldest_tensor.numel()
280
- )
281
- self.buffer_size -= oldest_tensor_size
282
- logger.debug(
283
- "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
284
- " buffer_size:%d, oldest_tensor_size:%d, rank:%d",
285
- remote_address,
286
- tensor_id,
287
- tensor_size,
288
- self.buffer_size,
289
- oldest_tensor_size,
290
- self.rank,
291
- )
292
-
293
- self.send_store[tensor_id] = tensor
294
- self.buffer_size += tensor_size
295
- logger.debug(
296
- "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
297
- "shape:%s, rank:%d, buffer_size:%d(%.2f%%)",
298
- remote_address,
299
- tensor_id,
300
- tensor_size,
301
- tensor.shape,
302
- self.rank,
303
- self.buffer_size,
304
- self.buffer_size / self.buffer_size_threshold * 100,
305
- )
306
- return True
307
-
308
- def recv_tensor(
309
- self,
310
- tensor_id: str,
311
- remote_address: str | None = None,
312
- ) -> torch.Tensor:
313
- if self.send_type == "PUT" or self.send_type == "PUT_ASYNC":
314
- start_time = time.time()
315
- with self.recv_store_cv:
316
- while tensor_id not in self.recv_store:
317
- self.recv_store_cv.wait()
318
- tensor = self.recv_store[tensor_id]
319
-
320
- if tensor is not None:
321
- if isinstance(tensor, tuple):
322
- addr, dtype, shape = tensor
323
- tensor = self.pool.load_tensor(addr, dtype, shape, self.device)
324
- else:
325
- self.buffer_size -= tensor.element_size() * tensor.numel()
326
- else:
327
- duration = time.time() - start_time
328
- logger.warning(
329
- "🔴[PUT]Recv From %s, tensor_id:%s, duration:%.3fms, rank:%d",
330
- remote_address,
331
- tensor_id,
332
- duration * 1000,
333
- self.rank,
334
- )
335
- return tensor
336
-
337
- # GET
338
- if remote_address is None:
339
- return None
340
-
341
- if remote_address not in self.socks:
342
- self.create_connect(remote_address)
343
-
344
- sock = self.socks[remote_address]
345
- comm, rank = self.comms[remote_address]
346
-
347
- data = {"cmd": "GET", "tensor_id": tensor_id}
348
- sock.send(msgpack.dumps(data))
349
-
350
- message = sock.recv()
351
- data = msgpack.loads(message)
352
- if data["ret"] != 0:
353
- logger.warning(
354
- "🔴[GET]Recv From %s, tensor_id: %s, ret: %d",
355
- remote_address,
356
- tensor_id,
357
- data["ret"],
358
- )
359
- return None
360
-
361
- with torch.cuda.stream(self.recv_stream):
362
- tensor = torch.empty(
363
- data["shape"], dtype=getattr(torch, data["dtype"]), device=self.device
364
- )
365
-
366
- self.recv(comm, tensor, rank ^ 1, self.recv_stream)
367
-
368
- return tensor
369
-
370
- def listen_for_requests(self):
371
- while True:
372
- socks = dict(self.poller.poll())
373
- if self.router_socket not in socks:
374
- continue
375
-
376
- remote_address, message = self.router_socket.recv_multipart()
377
- data = msgpack.loads(message)
378
- if data["cmd"] == "NEW":
379
- unique_id = self.nccl.unique_id_from_bytes(bytes(data["unique_id"]))
380
- with torch.accelerator.device_index(self.device.index):
381
- rank = 1
382
- with set_p2p_nccl_context(self.nccl_num_channels):
383
- comm: ncclComm_t = self.nccl.ncclCommInitRank(
384
- 2, unique_id, rank
385
- )
386
- self.comms[remote_address.decode()] = (comm, rank)
387
- logger.info(
388
- "🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
389
- self.zmq_address,
390
- remote_address.decode(),
391
- rank,
392
- )
393
- elif data["cmd"] == "PUT":
394
- tensor_id = data["tensor_id"]
395
- try:
396
- with torch.cuda.stream(self.recv_stream):
397
- tensor = torch.empty(
398
- data["shape"],
399
- dtype=getattr(torch, data["dtype"]),
400
- device=self.device,
401
- )
402
- self.router_socket.send_multipart([remote_address, b"0"])
403
- comm, rank = self.comms[remote_address.decode()]
404
- self.recv(comm, tensor, rank ^ 1, self.recv_stream)
405
- tensor_size = tensor.element_size() * tensor.numel()
406
- if self.buffer_size + tensor_size > self.buffer_size_threshold:
407
- # Store Tensor in memory pool
408
- addr = self.pool.store_tensor(tensor)
409
- tensor = (addr, tensor.dtype, tensor.shape)
410
- logger.warning(
411
- "🔴[PUT]Recv Tensor, Out Of Threshold, "
412
- "%s👈%s, data:%s, addr:%d",
413
- self.zmq_address,
414
- remote_address.decode(),
415
- data,
416
- addr,
417
- )
418
- else:
419
- self.buffer_size += tensor_size
420
-
421
- except torch.cuda.OutOfMemoryError:
422
- self.router_socket.send_multipart([remote_address, b"1"])
423
- tensor = None
424
- logger.warning(
425
- "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, data:%s",
426
- self.zmq_address,
427
- remote_address.decode(),
428
- data,
429
- )
430
-
431
- with self.recv_store_cv:
432
- self.recv_store[tensor_id] = tensor
433
- self.have_received_tensor_id(tensor_id)
434
- self.recv_store_cv.notify()
435
-
436
- elif data["cmd"] == "GET":
437
- tensor_id = data["tensor_id"]
438
- with self.send_store_cv:
439
- tensor = self.send_store.pop(tensor_id, None)
440
- if tensor is not None:
441
- data = {
442
- "ret": 0,
443
- "shape": tensor.shape,
444
- "dtype": str(tensor.dtype).replace("torch.", ""),
445
- }
446
- # LRU
447
- self.send_store[tensor_id] = tensor
448
- self.have_sent_tensor_id(tensor_id)
449
- else:
450
- data = {"ret": 1}
451
-
452
- self.router_socket.send_multipart([remote_address, msgpack.dumps(data)])
453
-
454
- if data["ret"] == 0:
455
- comm, rank = self.comms[remote_address.decode()]
456
- self.send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
457
- else:
458
- logger.warning(
459
- "🚧Unexpected, Received message from %s, data:%s",
460
- remote_address,
461
- data,
462
- )
463
-
464
- def have_sent_tensor_id(self, tensor_id: str):
465
- request_id = tensor_id.split("#")[0]
466
- if request_id not in self.send_request_id_to_tensor_ids:
467
- self.send_request_id_to_tensor_ids[request_id] = set()
468
- self.send_request_id_to_tensor_ids[request_id].add(tensor_id)
469
-
470
- def have_received_tensor_id(self, tensor_id: str):
471
- request_id = tensor_id.split("#")[0]
472
- if request_id not in self.recv_request_id_to_tensor_ids:
473
- self.recv_request_id_to_tensor_ids[request_id] = set()
474
- self.recv_request_id_to_tensor_ids[request_id].add(tensor_id)
475
-
476
- def send_async(self):
477
- while True:
478
- with self.send_queue_cv:
479
- while not self.send_queue:
480
- self.send_queue_cv.wait()
481
- item = self.send_queue.popleft()
482
- if not self.send_queue:
483
- self.send_queue_cv.notify()
484
- self.send_sync(item)
485
-
486
- def wait_for_sent(self):
487
- if self.send_type == "PUT_ASYNC":
488
- start_time = time.time()
489
- with self.send_queue_cv:
490
- while self.send_queue:
491
- self.send_queue_cv.wait()
492
- duration = time.time() - start_time
493
- logger.debug(
494
- "🚧[PUT_ASYNC]It took %.3fms to wait for the send_queue"
495
- " to be empty, rank:%d",
496
- duration * 1000,
497
- self.rank,
498
- )
499
-
500
- def send_sync(self, item: SendQueueItem) -> bool:
501
- if item.remote_address is None:
502
- return False
503
- if item.remote_address not in self.socks:
504
- self.create_connect(item.remote_address)
505
-
506
- tensor = item.tensor
507
-
508
- sock = self.socks[item.remote_address]
509
- comm, rank = self.comms[item.remote_address]
510
- data = {
511
- "cmd": "PUT",
512
- "tensor_id": item.tensor_id,
513
- "shape": tensor.shape,
514
- "dtype": str(tensor.dtype).replace("torch.", ""),
515
- }
516
- sock.send(msgpack.dumps(data))
517
-
518
- response = sock.recv()
519
- if response != b"0":
520
- logger.error(
521
- "🔴Send Tensor, Peer Out Of Memory/Threshold, %s 👉 %s, "
522
- "MyRank:%s, data:%s, tensor:%s, size:%fGB, response:%s",
523
- self.zmq_address,
524
- item.remote_address,
525
- rank,
526
- data,
527
- tensor.shape,
528
- tensor.element_size() * tensor.numel() / 1024**3,
529
- response.decode(),
530
- )
531
- return False
532
-
533
- self.send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
534
-
535
- if self.send_type == "PUT_ASYNC":
536
- self.have_sent_tensor_id(item.tensor_id)
537
-
538
- return True
539
-
540
- def get_finished(
541
- self, finished_req_ids: set[str], no_compile_layers
542
- ) -> tuple[set[str] | None, set[str] | None]:
543
- """
544
- Notifies worker-side connector ids of requests that have
545
- finished generating tokens.
546
-
547
- Returns:
548
- ids of requests that have finished asynchronous transfer,
549
- tuple of (sending/saving ids, recving/loading ids).
550
- The finished saves/sends req ids must belong to a set provided in a
551
- call to this method (this call or a prior one).
552
- """
553
-
554
- # Clear the buffer upon request completion.
555
- for request_id in finished_req_ids:
556
- for layer_name in no_compile_layers:
557
- tensor_id = request_id + "#" + layer_name
558
- if tensor_id in self.recv_store:
559
- with self.recv_store_cv:
560
- tensor = self.recv_store.pop(tensor_id, None)
561
- self.send_request_id_to_tensor_ids.pop(request_id, None)
562
- self.recv_request_id_to_tensor_ids.pop(request_id, None)
563
- if isinstance(tensor, tuple):
564
- addr, _, _ = tensor
565
- self.pool.free(addr)
566
-
567
- # TODO:Retrieve requests that have already sent the KV cache.
568
- finished_sending: set[str] = set()
569
-
570
- # TODO:Retrieve requests that have already received the KV cache.
571
- finished_recving: set[str] = set()
572
-
573
- return finished_sending or None, finished_recving or None
574
-
575
- def ping(self):
576
- sock = self.context.socket(zmq.DEALER)
577
- sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
578
- logger.debug("ping start, zmq_address:%s", self.zmq_address)
579
- sock.connect(f"tcp://{self.proxy_address}")
580
- data = {
581
- "type": "P" if self.config.is_kv_producer else "D",
582
- "http_address": self.http_address,
583
- "zmq_address": self.zmq_address,
584
- }
585
- while True:
586
- sock.send(msgpack.dumps(data))
587
- time.sleep(3)
588
-
589
- def send(self, comm, tensor: torch.Tensor, dst: int, stream=None):
590
- assert tensor.device == self.device, (
591
- f"this nccl communicator is created to work on {self.device}, "
592
- f"but the input tensor is on {tensor.device}"
593
- )
594
- if stream is None:
595
- stream = current_stream()
596
-
597
- with torch.cuda.stream(stream):
598
- self.nccl.ncclSend(
599
- buffer_type(tensor.data_ptr()),
600
- tensor.numel(),
601
- ncclDataTypeEnum.from_torch(tensor.dtype),
602
- dst,
603
- comm,
604
- cudaStream_t(stream.cuda_stream),
605
- )
606
- stream.synchronize()
607
-
608
- def recv(self, comm, tensor: torch.Tensor, src: int, stream=None):
609
- assert tensor.device == self.device, (
610
- f"this nccl communicator is created to work on {self.device}, "
611
- f"but the input tensor is on {tensor.device}"
612
- )
613
- if stream is None:
614
- stream = current_stream()
615
-
616
- with torch.cuda.stream(stream):
617
- self.nccl.ncclRecv(
618
- buffer_type(tensor.data_ptr()),
619
- tensor.numel(),
620
- ncclDataTypeEnum.from_torch(tensor.dtype),
621
- src,
622
- comm,
623
- cudaStream_t(stream.cuda_stream),
624
- )
625
- stream.synchronize()
626
-
627
- def close(self) -> None:
628
- self._listener_thread.join()
629
- if self.send_type == "PUT_ASYNC":
630
- self._send_thread.join()
631
- if self._ping_thread is not None:
632
- self._ping_thread.join()
vllm/distributed/kv_transfer/kv_connector/v1/p2p/tensor_memory_pool.py
deleted
-273
@@ -1,273 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import atexit
5
-import ctypes
6
-import math
7
-from dataclasses import dataclass
8
-
9
-import torch
10
-
11
-from vllm.logger import init_logger
12
-
13
-logger = init_logger(__name__)
14
-
15
-
16
-@dataclass
17
-class MemoryBlock:
18
- size: int
19
- addr: int
20
-
21
-
22
-"""A memory pool for managing pinned host memory allocations for tensors.
23
-
24
-This class implements a buddy allocation system to efficiently manage pinned
25
-host memory for tensor storage. It supports allocation, deallocation, and
26
-tensor storage/retrieval operations.
27
-
28
-Key Features:
29
-- Uses power-of-two block sizes for efficient buddy allocation
30
-- Supports splitting and merging of memory blocks
31
-- Provides methods to store CUDA tensors in pinned host memory
32
-- Allows loading tensors from pinned memory back to device
33
-- Automatically cleans up memory on destruction
34
-
35
-Attributes:
36
- max_block_size (int): Maximum block size (rounded to nearest power of two)
37
- min_block_size (int): Minimum block size (rounded to nearest power of two)
38
- free_lists (dict): Dictionary of free memory blocks by size
39
- allocated_blocks (dict): Dictionary of currently allocated blocks
40
- base_tensor (torch.Tensor): Base pinned memory tensor
41
- base_address (int): Base memory address of the pinned memory region
42
-
43
-Example:
44
- >>> pool = TensorMemoryPool(max_block_size=1024*1024)
45
- >>> tensor = torch.randn(100, device='cuda')
46
- >>> addr = pool.store_tensor(tensor)
47
- >>> loaded_tensor = pool.load_tensor(addr, tensor.dtype,
48
- ... tensor.shape, 'cuda')
49
- >>> pool.free(addr)
50
-"""
51
-
52
-
53
-class TensorMemoryPool:
54
- """Initializes the memory pool with given size constraints.
55
-
56
- Args:
57
- max_block_size (int): Maximum size of memory blocks to manage
58
- min_block_size (int, optional): Minimum size of memory blocks
59
- to manage. Defaults to 512.
60
-
61
- Raises:
62
- ValueError: If block sizes are invalid or max_block_size is less
63
- than min_block_size
64
- """
65
-
66
- def __init__(self, max_block_size: int, min_block_size: int = 512):
67
- if max_block_size <= 0 or min_block_size <= 0:
68
- raise ValueError("Block sizes must be positive")
69
- if max_block_size < min_block_size:
70
- raise ValueError("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(
94
- self.max_block_size // 4, dtype=torch.float32, pin_memory=True
95
- )
96
- self.base_address = self.base_tensor.data_ptr()
97
- initial_block = MemoryBlock(size=self.max_block_size, addr=self.base_address)
98
- self.free_lists[self.max_block_size][initial_block.addr] = initial_block
99
-
100
- logger.debug(
101
- "TensorMemoryPool, base_address:%d, max_block_size:%d",
102
- self.base_address,
103
- self.max_block_size,
104
- )
105
-
106
- def allocate(self, size: int) -> int:
107
- """Allocates a memory block of at least the requested size.
108
-
109
- Args:
110
- size (int): Minimum size of memory to allocate
111
-
112
- Returns:
113
- int: Address of the allocated memory block
114
-
115
- Raises:
116
- ValueError: If size is invalid or insufficient memory is available
117
- """
118
- if size <= 0:
119
- raise ValueError("Allocation size must be positive")
120
-
121
- required_size = self._round_to_power_of_two(max(size, self.min_block_size))
122
- if required_size > self.max_block_size:
123
- raise ValueError("Requested size exceeds maximum block size")
124
-
125
- current_size = required_size
126
- while current_size <= self.max_block_size:
127
- if self.free_lists[current_size]:
128
- _, block = self.free_lists[current_size].popitem()
129
- self._split_block(block, required_size)
130
- self.allocated_blocks[block.addr] = block
131
- return block.addr
132
- current_size *= 2
133
-
134
- raise ValueError("Insufficient memory")
135
-
136
- def _split_block(self, block: MemoryBlock, required_size: int):
137
- while block.size > required_size 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 = (
167
- block.size
168
- if (block.addr - self.base_address) % (2 * block.size) == 0
169
- else -block.size
170
- )
171
- buddy_addr = block.addr + buddy_offset
172
- buddy = self.free_lists[block.size].get(buddy_addr)
173
- if buddy:
174
- del self.free_lists[buddy.size][buddy.addr]
175
- merged_addr = min(block.addr, buddy.addr)
176
- merged_size = block.size * 2
177
- block = MemoryBlock(size=merged_size, addr=merged_addr)
178
- depth += 1
179
- else:
180
- break
181
- self.free_lists[block.size][block.addr] = block
182
-
183
- def store_tensor(self, tensor: torch.Tensor) -> int:
184
- """Stores a CUDA tensor in pinned host memory.
185
-
186
- Args:
187
- tensor (torch.Tensor): CUDA tensor to store
188
-
189
- Returns:
190
- int: Address where the tensor is stored
191
-
192
- Raises:
193
- ValueError: If tensor is not on CUDA or allocation fails
194
- """
195
- if not tensor.is_cuda:
196
- raise ValueError("Only CUDA tensors can be stored")
197
-
198
- size = tensor.element_size() * tensor.numel()
199
- addr = self.allocate(size)
200
- block = self.allocated_blocks[addr]
201
-
202
- if block.size < size:
203
- self.free(addr)
204
- raise ValueError(
205
- f"Allocated block size {block.size} is smaller than "
206
- f"required size {size}"
207
- )
208
-
209
- try:
210
- buffer = (ctypes.c_byte * block.size).from_address(block.addr)
211
- cpu_tensor = torch.frombuffer(
212
- buffer, dtype=tensor.dtype, count=tensor.numel()
213
- ).reshape(tensor.shape)
214
- except ValueError as err:
215
- self.free(addr)
216
- raise ValueError(f"Failed to create tensor view: {err}") from err
217
-
218
- cpu_tensor.copy_(tensor)
219
-
220
- return addr
221
-
222
- def load_tensor(
223
- self,
224
- addr: int,
225
- dtype: torch.dtype,
226
- shape: tuple[int, ...],
227
- device: torch.device,
228
- ) -> torch.Tensor:
229
- """Loads a tensor from pinned host memory to the specified device.
230
-
231
- Args:
232
- addr (int): Address where tensor is stored
233
- dtype (torch.dtype): Data type of the tensor
234
- shape (tuple[int, ...]): Shape of the tensor
235
- device: Target device for the loaded tensor
236
-
237
- Returns:
238
- torch.Tensor: The loaded tensor on the specified device
239
-
240
- Raises:
241
- ValueError: If address is invalid or sizes don't match
242
- """
243
- if addr not in self.allocated_blocks:
244
- raise ValueError("Invalid address to load")
245
-
246
- block = self.allocated_blocks[addr]
247
- num_elements = math.prod(shape)
248
- dtype_size = torch.tensor([], dtype=dtype).element_size()
249
- required_size = num_elements * dtype_size
250
-
251
- if required_size > block.size:
252
- raise ValueError("Requested tensor size exceeds block size")
253
-
254
- buffer = (ctypes.c_byte * block.size).from_address(block.addr)
255
- cpu_tensor = torch.frombuffer(buffer, dtype=dtype, count=num_elements).reshape(
256
- shape
257
- )
258
-
259
- cuda_tensor = torch.empty(shape, dtype=dtype, device=device)
260
-
261
- cuda_tensor.copy_(cpu_tensor)
262
-
263
- return cuda_tensor
264
-
265
- def cleanup(self):
266
- """Cleans up all memory resources and resets the pool state."""
267
- self.free_lists.clear()
268
- self.allocated_blocks.clear()
269
- if hasattr(self, "base_tensor"):
270
- del self.base_tensor
271
-
272
- def __del__(self):
273
- self.cleanup()