A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 21:00 UTC
[BugFix][PD]: make example proxy usable with P2pNcclConnector (#26628)
Signed-off-by: PAN <1162953505@qq.com>
Pan Li committed
Nov 21, 2025 at 01:38 UTC
e5bfcb6a88cda4f91e3c7074d7e76bb5d1d36362
2 files changed
+169
-99
benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py
+155
-94
@@ -5,11 +5,12 @@ 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
11
-from rate_limiter import RateLimiter
12
-from request_queue import RequestQueue
14
15
# Configure logging
16
logging.basicConfig(level=logging.INFO)
@@ -24,26 +25,8 @@ def parse_args():
25
parser.add_argument(
26
"--timeout",
27
type=float,
27
- default=300,
28
- help="Timeout for backend service requests in seconds (default: 300)",
29
- )
30
- parser.add_argument(
31
- "--max-concurrent",
32
- type=int,
33
- default=100,
34
- help="Maximum concurrent requests to backend services (default: 100)",
35
- )
36
- parser.add_argument(
37
- "--queue-size",
38
- type=int,
39
- default=500,
40
- help="Maximum number of requests in the queue (default: 500)",
41
- )
42
- parser.add_argument(
43
- "--rate-limit",
44
- type=int,
45
- default=40,
46
- help="Maximum requests per second (default: 40)",
28
+ default=6 * 60 * 60,
29
+ help="Timeout for backend service requests in seconds (default: 21600)",
30
)
31
parser.add_argument(
32
"--port",
@@ -54,14 +37,32 @@ def parse_args():
37
parser.add_argument(
38
"--prefill-url",
39
type=str,
57
- default="http://localhost:8100/v1/completions",
58
- help="Prefill service endpoint URL",
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,
63
- default="http://localhost:8200/v1/completions",
64
- help="Decode service endpoint URL",
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()
@@ -73,70 +74,129 @@ def main():
74
75
# Initialize configuration using command line parameters
76
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
76
- MAX_CONCURRENT_REQUESTS = args.max_concurrent
77
- REQUEST_QUEUE_SIZE = args.queue_size
78
- RATE_LIMIT = args.rate_limit
77
PREFILL_SERVICE_URL = args.prefill_url
78
DECODE_SERVICE_URL = args.decode_url
79
PORT = args.port
80
83
- app = Quart(__name__)
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
85
- # Initialize the rate limiter and request queue
86
- rate_limiter = RateLimiter(RATE_LIMIT)
87
- request_queue = RequestQueue(MAX_CONCURRENT_REQUESTS, REQUEST_QUEUE_SIZE)
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
89
- # Attach the configuration object to the application instance
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,
93
- "rate_limiter": rate_limiter,
94
- "request_queue": request_queue,
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
100
- # Start queue processing on app startup
101
- @app.before_serving
102
- async def startup():
103
- """Start request processing task when app starts serving"""
104
- asyncio.create_task(request_queue.process())
105
-
106
- async def forward_request(url, data):
107
- """Forward request to backend service with rate limiting and error handling"""
108
- headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
109
-
110
- # Use rate limiter as context manager
111
- async with (
112
- rate_limiter,
113
- aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
114
- ):
115
- try:
116
- async with session.post(
117
- url=url, json=data, headers=headers
118
- ) as response:
119
- if response.status == 200:
120
- # Stream response chunks
121
- async for chunk_bytes in response.content.iter_chunked(1024):
122
- yield chunk_bytes
123
- else:
124
- # Handle backend service errors
125
- error_text = await response.text()
126
- logger.error(
127
- "Backend service error: %s - %s",
128
- response.status,
129
- error_text,
130
- )
131
- yield b'{"error": "Backend service error"}'
132
- except aiohttp.ClientError as e:
133
- # Handle connection errors
134
- logger.error("Connection error to %s: %s", url, str(e))
135
- yield b'{"error": "Service unavailable"}'
136
- except asyncio.TimeoutError:
137
- # Handle timeout errors
138
- logger.error("Timeout connecting to %s", url)
139
- yield b'{"error": "Service timeout"}'
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"""
@@ -146,13 +206,27 @@ def main():
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
151
- async for _ in forward_request(PREFILL_SERVICE_URL, prefill_request):
152
- continue
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
155
- generator = forward_request(DECODE_SERVICE_URL, original_request_data)
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
@@ -168,23 +242,10 @@ def main():
242
@app.route("/v1/completions", methods=["POST"])
243
async def handle_request():
244
"""Handle incoming API requests with concurrency and rate limiting"""
171
- # Create task for request processing
172
- task = asyncio.create_task(process_request())
173
-
174
- # Enqueue request or reject if queue is full
175
- if not await request_queue.enqueue(task):
176
- return Response(
177
- response=b'{"error": "Server busy, try again later"}',
178
- status=503,
179
- content_type="application/json",
180
- )
181
-
245
try:
183
- # Return the response from the processing task
184
- return await task
246
+ return await process_request()
247
except asyncio.CancelledError:
186
- # Handle task cancellation (timeout or queue full)
187
- logger.warning("Request cancelled due to timeout or queue full")
248
+ logger.warning("Request cancelled")
249
return Response(
250
response=b'{"error": "Request cancelled"}',
251
status=503,
examples/online_serving/disaggregated_prefill.sh
+14
-5
@@ -24,7 +24,14 @@ cleanup() {
24
exit 0
25
}
26
27
-export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
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,7 +45,7 @@ fi
45
wait_for_server() {
46
local port=$1
47
timeout 1200 bash -c "
41
- until curl -s localhost:${port}/v1/completions > /dev/null; do
48
+ until curl -i localhost:${port}/v1/models > /dev/null; do
49
sleep 1
50
done" && return 0 || return 1
51
}
@@ -48,21 +55,23 @@ wait_for_server() {
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 \
56
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}' &
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
58
-# decoding instance, which is the KV consumer
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 \
65
- '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}' &
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