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

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 00:00 UTC

[P/D] rework mooncake connector and introduce its bootstrap server (#31034)

Signed-off-by: Tianchen Ding <dtcccc@linux.alibaba.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>

dtc committed Feb 4, 2026 at 00:08 UTC 0d6ccf68fa2c439e17d02f26c4044ed5df7f7099
9 files changed +1324 -203
docs/features/disagg_prefill.md
+2 -1
@@ -19,12 +19,13 @@ Two main reasons:
19
20 Please refer to [examples/online_serving/disaggregated_prefill.sh](../../examples/online_serving/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
21
22 -Now supports 5 types of connectors:
22 +Now supports 6 types of connectors:
23
24 - **ExampleConnector**: refer to [examples/offline_inference/disaggregated-prefill-v1/run.sh](../../examples/offline_inference/disaggregated-prefill-v1/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
25 - **LMCacheConnectorV1**: refer to [examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
26 - **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md).
27 - **P2pNcclConnector**: refer to [examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
28 +- **MooncakeConnector**: refer to [examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh) for the example usage of ExampleConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
29 - **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:
30
31 ```bash
docs/features/mooncake_connector_usage.md
+18 -7
@@ -31,11 +31,9 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
31 ### Proxy
32
33 ```bash
34 -python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --prefiller-host 192.168.0.2 --prefiller-port 8010 --decoder-host 192.168.0.3 --decoder-port 8020
34 +python examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
35 ```
36
37 -> NOTE: The Mooncake Connector currently uses the proxy from nixl_integration. This will be replaced with a self-developed proxy in the future.
38 -
37 Now you can send requests to the proxy server through port 8000.
38
39 ## Environment Variables
@@ -43,16 +41,29 @@ Now you can send requests to the proxy server through port 8000.
41 - `VLLM_MOONCAKE_BOOTSTRAP_PORT`: Port for Mooncake bootstrap server
42 - Default: 8998
43 - Required only for prefiller instances
46 - - Each vLLM worker needs a unique port on its host; using the same port number across different hosts is fine
47 - - For TP/DP deployments, each worker's port on a node is computed as: base_port + dp_rank * tp_size + tp_rank
48 - - Used for the decoder notifying the prefiller
44 + - For headless instances, must be the same as the master instance
45 + - Each instance needs a unique port on its host; using the same port number across different hosts is fine
46
47 - `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
48 - Default: 480
49 - If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
50
54 -## KV Role Options
51 +## KV Transfer Config
52 +
53 +### KV Role Options
54
55 - **kv_producer**: For prefiller instances that generate KV caches
56 - **kv_consumer**: For decoder instances that consume KV caches from prefiller
57 - **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined.
58 +
59 +### kv_connector_extra_config
60 +
61 +- **num_workers**: Size of thread pool for one prefiller worker to transfer KV caches by mooncake. (default 10)
62 +- **mooncake_protocol**: Mooncake connector protocol. (default "rdma")
63 +
64 +## Example Scripts/Code
65 +
66 +Refer to these example scripts in the vLLM repository:
67 +
68 +- [run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh)
69 +- [mooncake_connector_proxy.py](../../examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py)
examples/online_serving/disaggregated_serving/README.md
+1
@@ -6,3 +6,4 @@ This example contains scripts that demonstrate the disaggregated serving feature
6
7 - `disagg_proxy_demo.py` - Demonstrates XpYd (X prefill instances, Y decode instances).
8 - `kv_events.sh` - Demonstrates KV cache event publishing.
9 +- `mooncake_connector` - A proxy demo for MooncakeConnector.
examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py new
+376
@@ -0,0 +1,376 @@
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 ipaddress
7 +import itertools
8 +import os
9 +import urllib
10 +import uuid
11 +from contextlib import asynccontextmanager
12 +from typing import Any
13 +
14 +import httpx
15 +from fastapi import FastAPI, HTTPException, Request
16 +from fastapi.responses import StreamingResponse
17 +
18 +
19 +def maybe_wrap_ipv6_address(address: str) -> str:
20 + try:
21 + ipaddress.IPv6Address(address)
22 + return f"[{address}]"
23 + except ValueError:
24 + return address
25 +
26 +
27 +def make_http_path(host: str, port: int) -> str:
28 + return f"http://{host}:{port}"
29 +
30 +
31 +def prefiller_cycle(prefill_clients: list[Any]):
32 + while True:
33 + for prefill_client in prefill_clients:
34 + for i in range(prefill_client["dp_size"]):
35 + yield prefill_client, i
36 +
37 +
38 +async def get_prefiller_info(prefill_clients: list, ready: asyncio.Event):
39 + for prefill_client in prefill_clients:
40 + while True:
41 + try:
42 + # Wait for prefill service to be ready
43 + response = await prefill_client["client"].get("/health")
44 + response.raise_for_status()
45 + except Exception:
46 + await asyncio.sleep(1)
47 + continue
48 +
49 + response = await prefill_client["client"].get(
50 + prefill_client["bootstrap_addr"] + "/query"
51 + )
52 + response.raise_for_status()
53 + data = response.json()
54 + break
55 +
56 + for dp_rank, dp_entry in data.items():
57 + prefill_client["dp_engine_id"][int(dp_rank)] = dp_entry["engine_id"]
58 + dp_size = len(data)
59 + prefill_client["dp_size"] = dp_size
60 + print(f"Inited prefiller {prefill_client['url']} with dp_size={dp_size}")
61 +
62 + ready.set()
63 + print("All prefiller instances are ready.")
64 +
65 +
66 +@asynccontextmanager
67 +async def lifespan(app: FastAPI):
68 + """
69 + Lifespan context manager to handle startup and shutdown events.
70 + """
71 + # Startup: Initialize client pools for prefiller and decoder services
72 + app.state.prefill_clients = []
73 + app.state.decode_clients = []
74 + app.state.ready = asyncio.Event()
75 +
76 + # Create prefill clients
77 + for i, (url, bootstrap_port) in enumerate(global_args.prefill):
78 + parsed_url = urllib.parse.urlparse(url)
79 + hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
80 + app.state.prefill_clients.append(
81 + {
82 + "client": httpx.AsyncClient(
83 + timeout=None,
84 + base_url=url,
85 + limits=httpx.Limits(
86 + max_connections=None,
87 + max_keepalive_connections=None,
88 + ),
89 + ),
90 + "url": url,
91 + "bootstrap_addr": make_http_path(hostname, bootstrap_port or 8998),
92 + "dp_engine_id": {},
93 + }
94 + )
95 +
96 + # Create decode clients
97 + for i, url in enumerate(global_args.decode):
98 + parsed_url = urllib.parse.urlparse(url)
99 + hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
100 + app.state.decode_clients.append(
101 + {
102 + "client": httpx.AsyncClient(
103 + timeout=None,
104 + base_url=url,
105 + limits=httpx.Limits(
106 + max_connections=None,
107 + max_keepalive_connections=None,
108 + ),
109 + ),
110 + }
111 + )
112 +
113 + asyncio.create_task(get_prefiller_info(app.state.prefill_clients, app.state.ready))
114 +
115 + # Initialize round-robin iterators
116 + app.state.prefill_iterator = prefiller_cycle(app.state.prefill_clients)
117 + app.state.decode_iterator = itertools.cycle(range(len(app.state.decode_clients)))
118 +
119 + print(
120 + f"Got {len(app.state.prefill_clients)} prefill clients "
121 + f"and {len(app.state.decode_clients)} decode clients."
122 + )
123 +
124 + yield
125 +
126 + # Shutdown: Close all clients
127 + for client_info in app.state.prefill_clients:
128 + await client_info["client"].aclose()
129 +
130 + for client_info in app.state.decode_clients:
131 + await client_info["client"].aclose()
132 +
133 +
134 +# Update FastAPI app initialization to use lifespan
135 +app = FastAPI(lifespan=lifespan)
136 +
137 +
138 +def parse_args():
139 + parser = argparse.ArgumentParser()
140 +
141 + parser.add_argument("--port", type=int, default=8000)
142 + # Always use 127.0.0.1 as localhost binds to IPv6 which is blocked on CI
143 + parser.add_argument("--host", type=str, default="127.0.0.1")
144 +
145 + # For prefiller instances
146 + parser.add_argument(
147 + "--prefill",
148 + nargs="+",
149 + action="append",
150 + dest="prefill_raw",
151 + metavar=("URL", "bootstrap_port"),
152 + help=(
153 + "Prefill server URL and optional bootstrap port. "
154 + "Can be specified multiple times. "
155 + "Format: --prefill URL [BOOTSTRAP_PORT]. "
156 + "BOOTSTRAP_PORT can be a port number, "
157 + "'none', or omitted (defaults to none)."
158 + ),
159 + )
160 +
161 + # For decoder instances
162 + parser.add_argument(
163 + "--decode",
164 + nargs=1,
165 + action="append",
166 + dest="decode_raw",
167 + metavar=("URL",),
168 + help="Decode server URL. Can be specified multiple times.",
169 + )
170 +
171 + args = parser.parse_args()
172 + args.prefill = _parse_prefill_urls(args.prefill_raw)
173 + args.decode = _parse_decode_urls(args.decode_raw)
174 +
175 + return args
176 +
177 +
178 +# From sglang router_args.py
179 +def _parse_prefill_urls(prefill_list):
180 + """Parse prefill URLs from --prefill arguments.
181 +
182 + Format: --prefill URL [BOOTSTRAP_PORT]
183 + Example:
184 + --prefill http://prefill1:8080 9000 # With bootstrap port
185 + --prefill http://prefill2:8080 none # Explicitly no bootstrap port
186 + --prefill http://prefill3:8080 # Defaults to no bootstrap port
187 + """
188 + if not prefill_list:
189 + return []
190 +
191 + prefill_urls = []
192 + for prefill_args in prefill_list:
193 + url = prefill_args[0]
194 +
195 + # Handle optional bootstrap port
196 + if len(prefill_args) >= 2:
197 + bootstrap_port_str = prefill_args[1]
198 + # Handle 'none' as None
199 + if bootstrap_port_str.lower() == "none":
200 + bootstrap_port = None
201 + else:
202 + try:
203 + bootstrap_port = int(bootstrap_port_str)
204 + except ValueError as e:
205 + raise ValueError(
206 + f"Invalid bootstrap port: {bootstrap_port_str}. Must be a number or 'none'" # noqa: E501
207 + ) from e
208 + else:
209 + # No bootstrap port specified, default to None
210 + bootstrap_port = None
211 +
212 + prefill_urls.append((url, bootstrap_port))
213 +
214 + return prefill_urls
215 +
216 +
217 +def _parse_decode_urls(decode_list):
218 + """Parse decode URLs from --decode arguments.
219 +
220 + Format: --decode URL
221 + Example: --decode http://decode1:8081 --decode http://decode2:8081
222 + """
223 + if not decode_list:
224 + return []
225 +
226 + # decode_list is a list of single-element lists due to nargs=1
227 + return [url[0] for url in decode_list]
228 +
229 +
230 +def get_next_client(app, service_type: str):
231 + """
232 + Get the next client in round-robin fashion.
233 +
234 + Args:
235 + app: The FastAPI app instance
236 + service_type: Either 'prefill' or 'decode'
237 +
238 + Returns:
239 + The next client to use
240 + """
241 + if service_type == "prefill":
242 + return next(app.state.prefill_iterator)
243 + elif service_type == "decode":
244 + client_idx = next(app.state.decode_iterator)
245 + return app.state.decode_clients[client_idx]
246 + else:
247 + raise ValueError(f"Unknown service type: {service_type}")
248 +
249 +
250 +async def send_request_to_service(
251 + client_info: dict, dp_rank: int, endpoint: str, req_data: dict, request_id: str
252 +):
253 + """
254 + Send a request to a service using a client from the pool.
255 + """
256 + req_data = req_data.copy()
257 + req_data["kv_transfer_params"] = {
258 + "do_remote_decode": True,
259 + "do_remote_prefill": False,
260 + "transfer_id": f"xfer-{request_id}",
261 + }
262 + req_data["stream"] = False
263 + req_data["max_tokens"] = 1
264 + if "max_completion_tokens" in req_data:
265 + req_data["max_completion_tokens"] = 1
266 + if "stream_options" in req_data:
267 + del req_data["stream_options"]
268 + headers = {
269 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
270 + "X-Request-Id": request_id,
271 + "X-data-parallel-rank": str(dp_rank),
272 + }
273 +
274 + response = await client_info["client"].post(
275 + endpoint, json=req_data, headers=headers
276 + )
277 + response.raise_for_status()
278 +
279 + # CRITICAL: Release connection back to pool
280 + await response.aclose()
281 +
282 +
283 +async def stream_service_response(
284 + prefill_client_info: dict,
285 + prefill_dp_rank: int,
286 + decode_client_info: dict,
287 + endpoint: str,
288 + req_data: dict,
289 + request_id: str,
290 +):
291 + """
292 + Asynchronously stream response from a service using a client from the pool.
293 + """
294 + headers = {
295 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
296 + "X-Request-Id": request_id,
297 + }
298 +
299 + req_data["kv_transfer_params"] = {
300 + "do_remote_decode": False,
301 + "do_remote_prefill": True,
302 + "remote_bootstrap_addr": prefill_client_info["bootstrap_addr"],
303 + "remote_engine_id": prefill_client_info["dp_engine_id"][prefill_dp_rank],
304 + "transfer_id": f"xfer-{request_id}",
305 + }
306 +
307 + async with decode_client_info["client"].stream(
308 + "POST", endpoint, json=req_data, headers=headers
309 + ) as response:
310 + response.raise_for_status()
311 + async for chunk in response.aiter_bytes():
312 + yield chunk
313 +
314 +
315 +async def _handle_completions(api: str, request: Request):
316 + if not app.state.ready.is_set():
317 + raise HTTPException(status_code=503, detail="Service Unavailable")
318 +
319 + try:
320 + req_data = await request.json()
321 + request_id = str(uuid.uuid4())
322 +
323 + # Get the next prefill client in round-robin fashion
324 + prefill_client_info, prefill_dp_rank = get_next_client(request.app, "prefill")
325 +
326 + # Send request to prefill service
327 + asyncio.create_task(
328 + send_request_to_service(
329 + prefill_client_info, prefill_dp_rank, api, req_data, request_id
330 + )
331 + )
332 +
333 + decode_client_info = get_next_client(request.app, "decode")
334 +
335 + # Stream response from decode service
336 + async def generate_stream():
337 + async for chunk in stream_service_response(
338 + prefill_client_info,
339 + prefill_dp_rank,
340 + decode_client_info,
341 + api,
342 + req_data,
343 + request_id=request_id,
344 + ):
345 + yield chunk
346 +
347 + return StreamingResponse(generate_stream(), media_type="application/json")
348 +
349 + except Exception as e:
350 + import sys
351 + import traceback
352 +
353 + exc_info = sys.exc_info()
354 + print(f"Error occurred in disagg prefill proxy server - {api} endpoint")
355 + print(e)
356 + print("".join(traceback.format_exception(*exc_info)))
357 + raise
358 +
359 +
360 +@app.post("/v1/completions")
361 +async def handle_completions(request: Request):
362 + return await _handle_completions("/v1/completions", request)
363 +
364 +
365 +@app.post("/v1/chat/completions")
366 +async def handle_chat_completions(request: Request):
367 + return await _handle_completions("/v1/chat/completions", request)
368 +
369 +
370 +if __name__ == "__main__":
371 + global global_args
372 + global_args = parse_args()
373 +
374 + import uvicorn
375 +
376 + uvicorn.run(app, host=global_args.host, port=global_args.port)
examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh new
+222
@@ -0,0 +1,222 @@
1 +#!/bin/bash
2 +
3 +# =============================================================================
4 +# vLLM Disaggregated Serving Script for Mooncake Connector
5 +# =============================================================================
6 +# This script demonstrates disaggregated prefill and decode serving using
7 +# Mooncake Connector.
8 +#
9 +# Configuration can be customized via environment variables:
10 +# MODEL: Model to serve
11 +# PREFILL_GPUS: Comma-separated GPU IDs for prefill servers
12 +# DECODE_GPUS: Comma-separated GPU IDs for decode servers
13 +# PREFILL_PORTS: Comma-separated ports for prefill servers
14 +# BOOTSTRAP_PORTS: Bootstrap server port launched by prefill servers
15 +# DECODE_PORTS: Comma-separated ports for decode servers
16 +# PROXY_PORT: Proxy server port used to setup P/D disaggregated connection.
17 +# TIMEOUT_SECONDS: Server startup timeout
18 +# =============================================================================
19 +
20 +# Configuration - can be overridden via environment variables
21 +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct}
22 +TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200}
23 +PROXY_PORT=${PROXY_PORT:-8000}
24 +
25 +PREFILL_GPUS=${PREFILL_GPUS:-0}
26 +DECODE_GPUS=${DECODE_GPUS:-1}
27 +PREFILL_PORTS=${PREFILL_PORTS:-8010}
28 +BOOTSTRAP_PORTS=${BOOTSTRAP_PORTS:-8998}
29 +DECODE_PORTS=${DECODE_PORTS:-8020}
30 +
31 +echo "Warning: Mooncake Connector support for vLLM v1 is experimental and subject to change."
32 +echo ""
33 +echo "Architecture Configuration:"
34 +echo " Model: $MODEL"
35 +echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS, Bootstrap Port:$BOOTSTRAP_PORTS"
36 +echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS"
37 +echo " Proxy Port: $PROXY_PORT"
38 +echo " Timeout: ${TIMEOUT_SECONDS}s"
39 +echo ""
40 +
41 +PIDS=()
42 +
43 +# Switch to the directory of the current script
44 +cd "$(dirname "${BASH_SOURCE[0]}")"
45 +
46 +check_required_files() {
47 + local files=("mooncake_connector_proxy.py")
48 + for file in "${files[@]}"; do
49 + if [[ ! -f "$file" ]]; then
50 + echo "Required file $file not found in $(pwd)"
51 + exit 1
52 + fi
53 + done
54 +}
55 +
56 +check_hf_token() {
57 + if [ -z "$HF_TOKEN" ]; then
58 + echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
59 + echo "Example: export HF_TOKEN=your_token_here"
60 + exit 1
61 + fi
62 + if [[ "$HF_TOKEN" != hf_* ]]; then
63 + echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
64 + exit 1
65 + fi
66 + echo "HF_TOKEN is set and valid."
67 +}
68 +
69 +check_num_gpus() {
70 + # Check if the number of GPUs are >=2 via nvidia-smi
71 + num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
72 + if [ "$num_gpus" -lt 2 ]; then
73 + echo "You need at least 2 GPUs to run disaggregated prefill."
74 + exit 1
75 + else
76 + echo "Found $num_gpus GPUs."
77 + fi
78 +}
79 +
80 +ensure_python_library_installed() {
81 + echo "Checking if $1 is installed..."
82 + if ! python3 -c "import $1" > /dev/null 2>&1; then
83 + echo "$1 is not installed. Please install it via pip install $1."
84 + exit 1
85 + else
86 + echo "$1 is installed."
87 + fi
88 +}
89 +
90 +cleanup() {
91 + echo "Stopping everything…"
92 + trap - INT TERM # prevent re-entrancy
93 + pkill -9 -f "mooncake_connector_proxy.py"
94 + kill -- -$$ # negative PID == "this whole process-group"
95 + wait # reap children so we don't leave zombies
96 + exit 0
97 +}
98 +
99 +wait_for_server() {
100 + local port=$1
101 + local timeout_seconds=$TIMEOUT_SECONDS
102 + local start_time=$(date +%s)
103 +
104 + echo "Waiting for server on port $port..."
105 +
106 + while true; do
107 + if curl -s "localhost:${port}/v1/completions" > /dev/null; then
108 + echo "Server on port $port is ready."
109 + return 0
110 + fi
111 +
112 + local now=$(date +%s)
113 + if (( now - start_time >= timeout_seconds )); then
114 + echo "Timeout waiting for server on port $port"
115 + return 1
116 + fi
117 +
118 + sleep 1
119 + done
120 +}
121 +
122 +main() {
123 + check_required_files
124 + check_hf_token
125 + check_num_gpus
126 + ensure_python_library_installed vllm
127 + ensure_python_library_installed mooncake.engine
128 +
129 + trap cleanup INT
130 + trap cleanup USR1
131 + trap cleanup TERM
132 +
133 + echo "Launching disaggregated serving components..."
134 + echo "Please check the log files for detailed output:"
135 + echo " - prefill*.log: Prefill server logs"
136 + echo " - decode*.log: Decode server logs"
137 + echo " - proxy.log: Proxy server log"
138 +
139 + # Parse GPU and port arrays
140 + IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS"
141 + IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS"
142 + IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS"
143 + IFS=',' read -ra BOOTSTRAP_PORT_ARRAY <<< "$BOOTSTRAP_PORTS"
144 + IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
145 +
146 + proxy_param=""
147 +
148 + # =============================================================================
149 + # Launch Prefill Servers (X Producers)
150 + # =============================================================================
151 + echo ""
152 + echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..."
153 + for i in "${!PREFILL_GPU_ARRAY[@]}"; do
154 + local gpu_id=${PREFILL_GPU_ARRAY[$i]}
155 + local port=${PREFILL_PORT_ARRAY[$i]}
156 + local bootstrap_port=${BOOTSTRAP_PORT_ARRAY[$i]}
157 +
158 + echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, Bootstrap Port $bootstrap_port"
159 + VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
160 + --port $port \
161 + --kv-transfer-config \
162 + "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_producer\"}" > prefill$((i+1)).log 2>&1 &
163 + PIDS+=($!)
164 + proxy_param="${proxy_param} --prefill http://0.0.0.0:${port} $bootstrap_port"
165 + done
166 +
167 + # =============================================================================
168 + # Launch Decode Servers (Y Decoders)
169 + # =============================================================================
170 + echo ""
171 + echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..."
172 + for i in "${!DECODE_GPU_ARRAY[@]}"; do
173 + local gpu_id=${DECODE_GPU_ARRAY[$i]}
174 + local port=${DECODE_PORT_ARRAY[$i]}
175 +
176 + echo " Decode server $((i+1)): GPU $gpu_id, Port $port"
177 + CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
178 + --port $port \
179 + --kv-transfer-config \
180 + "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_consumer\"}" > decode$((i+1)).log 2>&1 &
181 + PIDS+=($!)
182 + proxy_param="${proxy_param} --decode http://0.0.0.0:${port}"
183 + done
184 +
185 + # =============================================================================
186 + # Launch Proxy Server
187 + # =============================================================================
188 + echo ""
189 + echo "Starting proxy server on port $PROXY_PORT..."
190 + python3 mooncake_connector_proxy.py $proxy_param --port $PROXY_PORT > proxy.log 2>&1 &
191 + PIDS+=($!)
192 +
193 + # =============================================================================
194 + # Wait for All Servers to Start
195 + # =============================================================================
196 + echo ""
197 + echo "Waiting for all servers to start..."
198 + for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
199 + if ! wait_for_server $port; then
200 + echo "Failed to start server on port $port"
201 + cleanup
202 + exit 1
203 + fi
204 + done
205 +
206 + echo ""
207 + echo "All servers are up. Starting benchmark..."
208 +
209 + # =============================================================================
210 + # Run Benchmark
211 + # =============================================================================
212 + vllm bench serve --port $PROXY_PORT --seed $(date +%s) \
213 + --backend vllm --model $MODEL \
214 + --dataset-name random --random-input-len 7500 --random-output-len 200 \
215 + --num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
216 +
217 + echo "Benchmarking done. Cleaning up..."
218 +
219 + cleanup
220 +}
221 +
222 +main
vllm/distributed/kv_transfer/kv_connector/factory.py
+1 -1
@@ -198,6 +198,6 @@ KVConnectorFactory.register_connector(
198 )
199 KVConnectorFactory.register_connector(
200 "MooncakeConnector",
201 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector",
201 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
202 "MooncakeConnector",
203 )
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py renamed
+577 -194
@@ -6,8 +6,10 @@ import time
6 from collections import defaultdict
7 from concurrent.futures import ThreadPoolExecutor
8 from dataclasses import dataclass
9 +from enum import IntEnum
10 from typing import TYPE_CHECKING, Any
11
12 +import httpx
13 import msgspec
14 import numpy as np
15 import torch
@@ -17,6 +19,7 @@ import zmq.asyncio
19 from vllm import envs
20 from vllm.config import VllmConfig
21 from vllm.distributed.kv_transfer.kv_connector.utils import (
22 + EngineId,
23 TpKVTopology,
24 get_current_attn_backend,
25 )
@@ -25,10 +28,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
28 KVConnectorMetadata,
29 KVConnectorRole,
30 )
31 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
32 + MooncakeBootstrapServer,
33 + RegisterWorkerPayload,
34 +)
35 from vllm.distributed.parallel_state import (
36 + get_pp_group,
37 get_tensor_model_parallel_rank,
38 get_tensor_model_parallel_world_size,
31 - get_tp_group,
39 + is_local_first_rank,
40 )
41 from vllm.forward_context import ForwardContext
42 from vllm.logger import init_logger
@@ -43,7 +51,7 @@ try:
51 except ImportError as e:
52 raise ImportError(
53 "Please install mooncake by following the instructions at "
46 - "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
54 + "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "
55 "to run VLLM with MooncakeTransferEngine."
56 ) from e
57
@@ -52,46 +60,75 @@ if TYPE_CHECKING:
60 from vllm.v1.kv_cache_interface import KVCacheConfig
61 from vllm.v1.request import Request
62
55 -EngineId = str
56 -ReqId = str
57 -
58 -TRANS_DONE = b"trans_done"
59 -TRANS_ERROR = b"trans_error"
63 +ReqId = str # Internal scheduler request ID
64 +TransferId = str # KV transfer coordination ID (shared by P/D)
65
66 logger = init_logger(__name__)
67
68
64 -class MooncakeAgentMetadata(
69 +class MooncakeXferMetadata(
70 msgspec.Struct,
71 omit_defaults=True, # type: ignore[call-arg]
67 - # required for @cached_property.
68 - dict=True,
72 ):
73 remote_hostname: str
74 remote_port: int
72 - request_ids: list[ReqId]
75 + remote_tp_size: int
76 + remote_tp_rank: int
77 + req_blocks: dict[ReqId, tuple[TransferId, list[int]]]
78 kv_caches_base_addr: list[int]
74 - block_ids: list[list[int]]
79 +
80 +
81 +class MooncakeXferResponseStatus(IntEnum):
82 + # Transfer finished
83 + FINISH = 0
84 + # Continue to receive
85 + CONTINUE = 1
86 + # Something wrong, see err_msg
87 + ERROR = 2
88 +
89 +
90 +class MooncakeXferResponse(
91 + msgspec.Struct,
92 + omit_defaults=True, # type: ignore[call-arg]
93 +):
94 + status: MooncakeXferResponseStatus
95 + ok_reqs: list[ReqId] | None = None
96 + err_reqs: list[ReqId] | None = None
97 + err_msg: str | None = None
98
99
100 @dataclass
78 -class RecvReqMeta:
101 +class PullReqMeta:
102 + d_req_id: ReqId
103 + transfer_id: TransferId
104 local_block_ids: list[int]
80 - remote_host: str
81 - remote_port: int
105 + remote_engine_id: EngineId
106 + remote_bootstrap_addr: str
107 + # Set expire time to avoid infinitely sending requests.
108 + expire_time: float = float("inf")
109 + # Designed for one D pairing to multiple P
110 + pull_tasks_count: int = 0
111
112
113 @dataclass
114 class SendBlockMeta:
115 + p_req_id: ReqId
116 + transfer_id: TransferId
117 local_block_ids: list[int]
118 ready: asyncio.Event
119 expire_time: float = float("inf")
120 + need_send: int = 0
121 + sent: int = 0
122 + sending: int = 0
123
124
125 class MooncakeConnectorMetadata(KVConnectorMetadata):
126 def __init__(self):
93 - self.reqs_to_recv: dict[ReqId, RecvReqMeta] = {}
94 - self.reqs_to_send: dict[ReqId, list[int]] = {}
127 + # Use (engine_id, dp_rank) to group reqs with same dp.
128 + # See comments in MooncakeBootstrapServer.
129 + self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict)
130 + self.reqs_to_send: dict[ReqId, tuple[TransferId, list[int]]] = {}
131 + self.reqs_not_processed: set[TransferId] = set()
132
133 def add_new_req(
134 self,
@@ -100,14 +137,18 @@ class MooncakeConnectorMetadata(KVConnectorMetadata):
137 kv_transfer_params: dict[str, Any],
138 load_remote_cache: bool = True,
139 ):
140 + transfer_id = kv_transfer_params["transfer_id"]
141 if load_remote_cache:
104 - self.reqs_to_recv[request_id] = RecvReqMeta(
142 + remote_engine_id = kv_transfer_params["remote_engine_id"]
143 + self.reqs_to_recv[remote_engine_id][request_id] = PullReqMeta(
144 + d_req_id=request_id,
145 local_block_ids=local_block_ids,
106 - remote_host=kv_transfer_params["remote_host"],
107 - remote_port=kv_transfer_params["remote_port"],
146 + remote_engine_id=remote_engine_id,
147 + remote_bootstrap_addr=kv_transfer_params["remote_bootstrap_addr"],
148 + transfer_id=transfer_id,
149 )
150 else:
110 - self.reqs_to_send[request_id] = local_block_ids
151 + self.reqs_to_send[request_id] = (transfer_id, local_block_ids)
152
153
154 class MooncakeConnector(KVConnectorBase_V1):
@@ -209,19 +250,24 @@ class MooncakeConnectorScheduler:
250
251 def __init__(self, vllm_config: VllmConfig, engine_id: str):
252 self.vllm_config = vllm_config
212 - self.engine_id: EngineId = engine_id
213 - self.side_channel_host = get_ip()
214 - self.side_channel_port = get_mooncake_side_channel_port(vllm_config)
253
254 assert vllm_config.kv_transfer_config
217 - self.kv_role = vllm_config.kv_transfer_config.kv_role
255 + self.is_kv_producer: bool = (
256 + vllm_config.kv_transfer_config.kv_role == "kv_producer"
257 + )
258 + self.is_kv_consumer: bool = (
259 + vllm_config.kv_transfer_config.kv_role == "kv_consumer"
260 + )
261 logger.info("Initializing Mooncake Transfer Engine Scheduler %s", engine_id)
262
263 # Requests that need to start recv/send.
264 # New requests are added by update_state_after_alloc in
265 # the scheduler. Used to make metadata passed to Worker.
266 self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
224 - self._reqs_need_send: dict[ReqId, list[int]] = {}
267 + self._reqs_need_send: dict[ReqId, tuple[Request, list[int]]] = {}
268 + # Reqs to remove from processed set because they're not to send after
269 + # remote prefill or aborted.
270 + self._reqs_not_processed: set[TransferId] = set()
271
272 def get_num_new_matched_tokens(
273 self, request: "Request", num_computed_tokens: int
@@ -249,8 +295,12 @@ class MooncakeConnectorScheduler:
295 params,
296 )
297
252 - if params is not None and params.get("do_remote_prefill"):
298 + if not params:
299 + return 0, False
300 +
301 + if params.get("do_remote_prefill"):
302 # Remote prefill: get all prompt blocks from remote.
303 + assert not self.is_kv_producer
304 token_ids = request.prompt_token_ids or []
305 count = len(token_ids) - num_computed_tokens
306 if count > 0:
@@ -265,7 +315,8 @@ class MooncakeConnectorScheduler:
315 params = request.kv_transfer_params
316 logger.debug(
317 "MooncakeConnector update_state_after_alloc: "
268 - "num_external_tokens=%s, kv_transfer_params=%s",
318 + "req_id=%s num_external_tokens=%s, kv_transfer_params=%s",
319 + request.request_id,
320 num_external_tokens,
321 params,
322 )
@@ -274,8 +325,11 @@ class MooncakeConnectorScheduler:
325 return
326
327 if params.get("do_remote_prefill"):
277 - assert self.kv_role != "kv_producer"
278 - if all(p in params for p in ("remote_host", "remote_port")):
328 + assert not self.is_kv_producer
329 + if all(
330 + p in params
331 + for p in ("remote_engine_id", "remote_bootstrap_addr", "transfer_id")
332 + ):
333 # If remote_blocks and num_external_tokens = 0, we have
334 # a full prefix cache hit on the D worker. We need to call
335 # send_notif in _read_blocks to free the memory on the P.
@@ -294,8 +348,12 @@ class MooncakeConnectorScheduler:
348 params["do_remote_prefill"] = False
349
350 elif params.get("do_remote_decode"):
297 - # Add an empty list to worker to create event.
298 - self._reqs_need_send[request.request_id] = []
351 + assert not self.is_kv_consumer
352 + if not params.get("transfer_id"):
353 + logger.warning("Missing transfer_id in kv_transfer_params from router!")
354 + else:
355 + # Add an empty list to worker to create event.
356 + self._reqs_need_send[request.request_id] = (request, [])
357
358 def build_connector_meta(
359 self,
@@ -303,8 +361,8 @@ class MooncakeConnectorScheduler:
361 ) -> KVConnectorMetadata:
362 meta = MooncakeConnectorMetadata()
363
306 - # Loop through scheduled reqs and convert to RecvReqMeta.
307 - if self.kv_role != "kv_producer":
364 + # Loop through scheduled reqs and convert to PullReqMeta.
365 + if not self.is_kv_producer:
366 for req_id, (req, block_ids) in self._reqs_need_recv.items():
367 assert req.kv_transfer_params is not None
368 meta.add_new_req(
@@ -314,15 +372,18 @@ class MooncakeConnectorScheduler:
372 )
373 self._reqs_need_recv.clear()
374
317 - if self.kv_role != "kv_consumer":
318 - for req_id, block_ids in self._reqs_need_send.items():
375 + if not self.is_kv_consumer:
376 + for req_id, (req, block_ids) in self._reqs_need_send.items():
377 + assert req.kv_transfer_params is not None
378 meta.add_new_req(
379 request_id=req_id,
380 local_block_ids=block_ids,
322 - kv_transfer_params={},
381 + kv_transfer_params=req.kv_transfer_params,
382 load_remote_cache=False,
383 )
384 self._reqs_need_send.clear()
385 + meta.reqs_not_processed = self._reqs_not_processed
386 + self._reqs_not_processed = set()
387
388 return meta
389
@@ -338,12 +399,13 @@ class MooncakeConnectorScheduler:
399
400 params = request.kv_transfer_params
401 logger.debug(
341 - "MooncakeConnector request_finished, request_status=%s, "
402 + "MooncakeConnector request_finished, req_id=%s, request_status=%s, "
403 "kv_transfer_params=%s",
404 + request.request_id,
405 request.status,
406 params,
407 )
346 - if not params:
408 + if not params or not params.get("transfer_id"):
409 return False, None
410
411 if params.get("do_remote_prefill"):
@@ -353,32 +415,30 @@ class MooncakeConnectorScheduler:
415 # To avoid stranding the prefill blocks in the prefill instance,
416 # we must add empty block_ids to _reqs_need_recv so that our
417 # worker side will notify and free blocks in the prefill instance.
356 - assert self.kv_role != "kv_producer"
418 + assert not self.is_kv_producer
419 self._reqs_need_recv[request.request_id] = (request, [])
420 params["do_remote_prefill"] = False
421 return False, None
422
361 - if (
362 - not params.get("do_remote_decode")
363 - or request.status != RequestStatus.FINISHED_LENGTH_CAPPED
364 - ):
423 + if not params.get("do_remote_decode"):
424 return False, None
425
367 - assert self.kv_role != "kv_consumer"
426 + assert not self.is_kv_consumer
427 +
428 + if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
429 + # Also include the case of a P/D Prefill request with immediate
430 + # block free (eg abort). Stop tracking this request.
431 + self._reqs_not_processed.add(params["transfer_id"])
432 + return False, None
433
434 # TODO: check whether block_ids actually ever be 0. If not we could
435 # remove the conditional below
436 delay_free_blocks = len(block_ids) > 0
437
438 if delay_free_blocks:
374 - self._reqs_need_send[request.request_id] = block_ids
439 + self._reqs_need_send[request.request_id] = (request, block_ids)
440
376 - return delay_free_blocks, dict(
377 - do_remote_prefill=True,
378 - do_remote_decode=False,
379 - remote_host=self.side_channel_host,
380 - remote_port=self.side_channel_port,
381 - )
441 + return delay_free_blocks, None
442
443
444 class MooncakeConnectorWorker:
@@ -391,7 +451,18 @@ class MooncakeConnectorWorker:
451
452 self.engine = TransferEngine()
453 self.hostname = get_ip()
394 - protocol = self.vllm_config.kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
454 +
455 + assert (kv_transfer_config := vllm_config.kv_transfer_config)
456 + self.is_kv_producer: bool = kv_transfer_config.kv_role == "kv_producer"
457 + self.is_kv_consumer: bool = kv_transfer_config.kv_role == "kv_consumer"
458 + self.num_sender_workers = kv_transfer_config.kv_connector_extra_config.get(
459 + "num_workers", 10
460 + )
461 + # Create more tasks than workers to keep the thread pool saturated.
462 + # Tasks can await async events, so a surplus (2x is a robust heuristic)
463 + # prevents workers from idling.
464 + self.num_sender_tasks = self.num_sender_workers * 2
465 + protocol = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
466 "mooncake_protocol", "rdma"
467 )
468 logger.info(
@@ -409,33 +480,31 @@ class MooncakeConnectorWorker:
480 self.rpc_port,
481 )
482
412 - # Mooncake handshake port.
413 - self.side_channel_port: int = get_mooncake_side_channel_port(vllm_config)
414 -
483 + self._remote_agents: dict[EngineId, dict[int, dict[int, str]]] = {}
484 + self._pending_bootstrap_querys: dict[str, asyncio.Event] = {}
485 + self.side_channel_port: int = 0 # we will bind it in register_kv_caches()
486 self.engine_id: EngineId = engine_id
487 self.tp_rank = get_tensor_model_parallel_rank()
417 - self.world_size = get_tensor_model_parallel_world_size()
418 - self.tp_group = get_tp_group()
488 + self.tp_size = get_tensor_model_parallel_world_size()
489 self.num_blocks = 0
490
421 - assert vllm_config.kv_transfer_config
422 - self.kv_role = vllm_config.kv_transfer_config.kv_role
423 - self.num_sender_workers = (
424 - vllm_config.kv_transfer_config.kv_connector_extra_config.get(
425 - "num_workers", 10
491 + assert (parallel_config := vllm_config.parallel_config)
492 + dp_rank = parallel_config.data_parallel_index
493 + dp_local_rank = parallel_config.data_parallel_rank_local
494 + self.dp_rank = dp_local_rank if parallel_config.local_engines_only else dp_rank
495 + pp_size = vllm_config.parallel_config.pipeline_parallel_size
496 + if pp_size > 1:
497 + raise ValueError(
498 + "Mooncake Transfer Engine does not support pipeline parallelism yet."
499 )
427 - )
428 - # Create more tasks than workers to keep the thread pool saturated.
429 - # Tasks can await async events, so a surplus (2x is a robust heuristic)
430 - # prevents workers from idling.
431 - self.num_sender_tasks = self.num_sender_workers * 2
500 + self.pp_rank = get_pp_group().rank_in_group
501
502 self.kv_caches_base_addr: list[int] = []
503 self.device_kv_caches: dict[str, torch.Tensor] = {}
435 - self.reqs_need_send: dict[ReqId, SendBlockMeta] = {}
504 + self.reqs_need_send: dict[TransferId, SendBlockMeta] = {}
505
506 # For kv_both, we will act both prefiller and decoder.
438 - if self.kv_role != "kv_consumer":
507 + if not self.is_kv_consumer:
508 # Background threads for sending kvcaches to D.
509 self._sender_executor = ThreadPoolExecutor(
510 max_workers=self.num_sender_workers,
@@ -454,7 +523,15 @@ class MooncakeConnectorWorker:
523 )
524 self._sender_listener_t.start()
525
457 - if self.kv_role != "kv_producer":
526 + # Start bootstrap server on global rank 0.
527 + if should_launch_bootstrap_server(vllm_config):
528 + _, port = get_mooncake_bootstrap_addr(vllm_config)
529 + self.bootstrap_server = MooncakeBootstrapServer(
530 + vllm_config, "0.0.0.0", port
531 + )
532 + self.bootstrap_server.start()
533 +
534 + if not self.is_kv_producer:
535 self.receiver_loop = asyncio.new_event_loop()
536 self._mooncake_receiver_t = threading.Thread(
537 target=_async_loop, args=(self.receiver_loop,), daemon=True
@@ -478,7 +555,7 @@ class MooncakeConnectorWorker:
555 logger.debug("Detected attention backend %s", self.backend_name)
556 logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
557
481 - self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
558 + self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
559 self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
560 self.kv_topo = TpKVTopology(
561 tp_rank=self.tp_rank,
@@ -492,7 +569,8 @@ class MooncakeConnectorWorker:
569
570 self.async_zmq_ctx = zmq.asyncio.Context()
571 self._encoder = msgspec.msgpack.Encoder()
495 - self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata)
572 + self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata)
573 + self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse)
574
575 def __del__(self):
576 self.shutdown()
@@ -500,26 +578,62 @@ class MooncakeConnectorWorker:
578 def shutdown(self):
579 """Cleanup background threads on destruction."""
580 self.async_zmq_ctx.term()
503 - if self.kv_role != "kv_consumer":
581 + if not self.is_kv_consumer:
582 self._sender_executor.shutdown(wait=False)
583 if self.sender_loop.is_running():
584 self.sender_loop.call_soon_threadsafe(self.sender_loop.stop)
585 self._sender_listener_t.join()
508 - if self.kv_role != "kv_producer" and self.receiver_loop.is_running():
586 + if should_launch_bootstrap_server(self.vllm_config):
587 + self.bootstrap_server.shutdown()
588 + if not self.is_kv_producer and self.receiver_loop.is_running():
589 self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop)
590 self._mooncake_receiver_t.join()
591
512 - async def _mooncake_sender_listener(
513 - self, ready_event: threading.Event, base_port: int, tp_rank: int
514 - ):
592 + async def register_worker_with_bootstrap(self):
593 + host, port = get_mooncake_bootstrap_addr(self.vllm_config)
594 + url = make_zmq_path("http", host, port) + "/register"
595 + worker_addr = make_zmq_path("tcp", self.hostname, self.side_channel_port)
596 + payload = RegisterWorkerPayload(
597 + engine_id=self.engine_id,
598 + dp_rank=self.dp_rank,
599 + tp_rank=self.tp_rank,
600 + pp_rank=self.pp_rank,
601 + addr=worker_addr,
602 + )
603 + while True:
604 + try:
605 + async with httpx.AsyncClient() as client:
606 + response = await client.post(url, json=payload.model_dump())
607 + response.raise_for_status()
608 + logger.debug("Successfully registered with bootstrap server at %s", url)
609 + break
610 + except httpx.ConnectError:
611 + # Bootstrap server not ready, wait for a while and retry.
612 + await asyncio.sleep(1)
613 + except Exception as e:
614 + err_msg = (
615 + e.response.text if isinstance(e, httpx.HTTPStatusError) else str(e)
616 + )
617 + logger.error(
618 + "Error registering %s with bootstrap server: %s", payload, err_msg
619 + )
620 + raise e
621 +
622 + async def _mooncake_sender_listener(self, ready_event: threading.Event):
623 """
624 Background thread that listens for Mooncake requests, dispatches them
625 to a thread pool, and sends acknowledgments upon completion.
626 """
627
520 - path = make_zmq_path("tcp", self.hostname, base_port + tp_rank)
521 - sock = make_zmq_socket(self.async_zmq_ctx, path, zmq.ROUTER)
522 - logger.debug("Mooncake sender starting listening on path: %s", path)
628 + sock = self.async_zmq_ctx.socket(zmq.ROUTER)
629 + self.side_channel_port = sock.bind_to_random_port(f"tcp://{self.hostname}")
630 + logger.debug(
631 + "Mooncake sender starting listening on path: tcp://%s:%d",
632 + self.hostname,
633 + self.side_channel_port,
634 + )
635 +
636 + await self.register_worker_with_bootstrap()
637
638 # Create async worker tasks that process items from the queue
639 sender_tasks = [
@@ -531,7 +645,7 @@ class MooncakeConnectorWorker:
645
646 try:
647 while True:
534 - identity, _, metadata_bytes = await sock.recv_multipart()
648 + identity, metadata_bytes = await sock.recv_multipart()
649 await self.sender_worker_queue.put((identity, metadata_bytes))
650 except zmq.ContextTerminated:
651 logger.debug("ZMQ context terminated, exiting Mooncake sender thread.")
@@ -549,12 +663,16 @@ class MooncakeConnectorWorker:
663 try:
664 identity, metadata_bytes = await self.sender_worker_queue.get()
665 try:
552 - metadata = self._decoder.decode(metadata_bytes)
553 - await self.send_kv_to_decode(metadata)
554 - await sock.send_multipart((identity, b"", TRANS_DONE))
666 + metadata = self._xfer_meta_decoder.decode(metadata_bytes)
667 + await self.send_kv_to_decode(identity, sock, metadata)
668 except Exception as e:
669 logger.error("Error processing Mooncake xfer request: %s", e)
557 - await sock.send_multipart((identity, b"", TRANS_ERROR))
670 + error_response = MooncakeXferResponse(
671 + status=MooncakeXferResponseStatus.ERROR, err_msg=str(e)
672 + )
673 + await sock.send_multipart(
674 + (identity, self._encoder.encode(error_response))
675 + )
676 finally:
677 self.sender_worker_queue.task_done()
678 except asyncio.CancelledError:
@@ -562,55 +680,169 @@ class MooncakeConnectorWorker:
680 except Exception as e:
681 logger.error("Error in _sender_worker: %s", e)
682
565 - async def send_kv_to_decode(self, meta: MooncakeAgentMetadata):
566 - send_reqs: list[tuple[ReqId, SendBlockMeta]] = []
567 - for req_id in meta.request_ids:
568 - send_meta = self.reqs_need_send.get(req_id)
569 - if send_meta is None:
570 - logger.warning("Request %s not found in reqs_need_send", req_id)
571 - return
572 - # Mark it as not expired. We will send it now.
573 - send_meta.expire_time = float("inf")
574 - send_reqs.append((req_id, send_meta))
575 -
576 - src_ptrs, dst_ptrs, lengths = await self._build_transfer_params(send_reqs, meta)
577 - remote_session = f"{meta.remote_hostname}:{meta.remote_port}"
578 - ret_value = await self.sender_loop.run_in_executor(
579 - self._sender_executor,
580 - self._send_blocks,
581 - remote_session,
582 - src_ptrs,
583 - dst_ptrs,
584 - lengths,
585 - )
683 + async def send_kv_to_decode(
684 + self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata
685 + ):
686 + pending_reqs: dict[ReqId, SendBlockMeta] = {}
687 + remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size)
688 + if self.tp_rank not in remote_tp_ranks:
689 + # This D worker does not pair with the P worker.
690 + msg = f"This P tp_rank {self.tp_rank} not in remote D target ranks {remote_tp_ranks}" # noqa: E501
691 + logger.error(msg)
692 + response = MooncakeXferResponse(
693 + status=MooncakeXferResponseStatus.ERROR,
694 + err_msg=msg,
695 + )
696 + await sock.send_multipart((identity, self._encoder.encode(response)))
697 + return
698 + for d_req_id, (transfer_id, _) in meta.req_blocks.items():
699 + if transfer_id not in self.reqs_need_send:
700 + # This req is not enqueued in P side yet, create it here.
701 + self.reqs_need_send[transfer_id] = SendBlockMeta(
702 + p_req_id="",
703 + transfer_id=transfer_id,
704 + local_block_ids=[],
705 + ready=asyncio.Event(),
706 + )
707 + send_meta = self.reqs_need_send[transfer_id]
708 + pending_reqs[d_req_id] = send_meta
709
587 - if ret_value != 0:
588 - raise RuntimeError(f"Error in batch_transfer_sync_write: {ret_value}")
710 + async def wait_and_ret(
711 + d_req_id: ReqId, send_meta: SendBlockMeta
712 + ) -> tuple[ReqId, SendBlockMeta]:
713 + await send_meta.ready.wait()
714 + return d_req_id, send_meta
715 +
716 + wait_tasks = [
717 + asyncio.create_task(wait_and_ret(d_req_id, send_meta))
718 + for d_req_id, send_meta in pending_reqs.items()
719 + ]
720 +
721 + while wait_tasks:
722 + done, pending = await asyncio.wait(
723 + wait_tasks,
724 + timeout=envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
725 + return_when=asyncio.FIRST_COMPLETED,
726 + )
727 +
728 + if not done:
729 + # Timeout, abort all pending requests.
730 + for task in wait_tasks:
731 + task.cancel()
732 + logger.warning(
733 + "Timeout waiting for P side ready: %s", list(pending_reqs)
734 + )
735 + response = MooncakeXferResponse(
736 + status=MooncakeXferResponseStatus.FINISH,
737 + err_reqs=list(pending_reqs),
738 + err_msg="Timeout waiting for P side ready.",
739 + )
740 + await sock.send_multipart((identity, self._encoder.encode(response)))
741 + break
742 +
743 + wait_tasks = list(pending)
744 + response_status = (
745 + MooncakeXferResponseStatus.CONTINUE
746 + if wait_tasks
747 + else MooncakeXferResponseStatus.FINISH
748 + )
749 + ready_reqs: list[tuple[ReqId, SendBlockMeta]] = []
750 + for task in done:
751 + d_req_id, send_meta = task.result()
752 + del pending_reqs[d_req_id]
753 + # Do we still in reqs_need_send (not expired)?
754 + if send_meta.transfer_id in self.reqs_need_send:
755 + # Mark it sending to avoid expiration.
756 + send_meta.sending += 1
757 + if not send_meta.need_send:
758 + self.resolve_need_send(send_meta, remote_tp_ranks)
759 + ready_reqs.append((d_req_id, send_meta))
760 + else:
761 + # Otherwise (expired, very unlikely), just forget it.
762 + logger.warning(
763 + "Request %s expired before sending on P side.", d_req_id
764 + )
765
590 - for req_id in meta.request_ids:
591 - del self.reqs_need_send[req_id]
766 + src_ptrs, dst_ptrs, lengths, err_reqs = await self._build_transfer_params(
767 + ready_reqs, meta
768 + )
769
593 - self.finished_sending_reqs.update(meta.request_ids)
770 + if err_reqs:
771 + response = MooncakeXferResponse(
772 + status=response_status,
773 + err_reqs=err_reqs,
774 + err_msg="P num blocks less than D",
775 + )
776 + await sock.send_multipart((identity, self._encoder.encode(response)))
777 +
778 + if src_ptrs:
779 + remote_session = f"{meta.remote_hostname}:{meta.remote_port}"
780 + ret_value = await self.sender_loop.run_in_executor(
781 + self._sender_executor,
782 + self._send_blocks,
783 + remote_session,
784 + src_ptrs,
785 + dst_ptrs,
786 + lengths,
787 + )
788 +
789 + if ret_value != 0:
790 + err_reqs = []
791 + for d_req_id, send_meta in ready_reqs:
792 + send_meta.sending -= 1
793 + err_reqs.append(d_req_id)
794 + # Do best effort to transfer the remaining reqs.
795 + response = MooncakeXferResponse(
796 + status=response_status,
797 + err_reqs=err_reqs,
798 + err_msg=f"Mooncake transfer engine returned {ret_value}",
799 + )
800 + await sock.send_multipart(
801 + (identity, self._encoder.encode(response))
802 + )
803 + continue
804 +
805 + for d_req_id, send_meta in ready_reqs:
806 + # TODO: for heterogeneous TP (one P pairs to multiple D),
807 + # we need to check whether all headers are sent.
808 + # If not, we should set expire_time to normal and skip the below.
809 + send_meta.sending -= 1
810 + send_meta.sent += 1
811 + if send_meta.sent == send_meta.need_send:
812 + del self.reqs_need_send[send_meta.transfer_id]
813 + self.finished_sending_reqs.add(send_meta.p_req_id)
814 +
815 + response = MooncakeXferResponse(
816 + status=response_status,
817 + ok_reqs=[d_req_id for d_req_id, _ in ready_reqs],
818 + )
819 + await sock.send_multipart((identity, self._encoder.encode(response)))
820 +
821 + def resolve_need_send(self, send_meta: SendBlockMeta, remote_tp_ranks: list[int]):
822 + # Prepare for heterogeneous TP (one P pairs to multiple D)
823 + send_meta.need_send = len(remote_tp_ranks)
824 + if send_meta.need_send != 1:
825 + logger.error("Mooncake: Heterogeneous TP is not supported yet.")
826 + raise NotImplementedError(
827 + "Mooncake: Heterogeneous TP is not supported yet."
828 + )
829
830 async def _build_transfer_params(
831 self,
597 - send_reqs: list[tuple[ReqId, SendBlockMeta]],
598 - agent_meta: MooncakeAgentMetadata,
599 - ) -> tuple[list[int], list[int], list[int]]:
832 + ready_reqs: list[tuple[ReqId, SendBlockMeta]],
833 + agent_meta: MooncakeXferMetadata,
834 + ) -> tuple[list[int], list[int], list[int], list[ReqId]]:
835 src_ptrs = []
836 dst_ptrs = []
837 lengths = []
838 + err_reqs: list[ReqId] = []
839 local_base_addr = self.kv_caches_base_addr
840 remote_base_addr = agent_meta.kv_caches_base_addr
841 block_len = self.block_len
842 remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}"
843
608 - assert len(send_reqs) == len(agent_meta.block_ids)
609 - for (req_id, send_meta), remote_block_ids in zip(
610 - send_reqs, agent_meta.block_ids
611 - ):
612 - await send_meta.ready.wait()
613 -
844 + for d_req_id, send_meta in ready_reqs:
845 + _, remote_block_ids = agent_meta.req_blocks[d_req_id]
846 num_remote_blocks = len(remote_block_ids)
847 if num_remote_blocks == 0:
848 continue
@@ -618,7 +850,15 @@ class MooncakeConnectorWorker:
850 local_block_ids = send_meta.local_block_ids
851 # Partial prefix cache hit: just read uncomputed blocks.
852 num_local_blocks = len(local_block_ids)
621 - assert num_local_blocks >= num_remote_blocks
853 + if num_local_blocks < num_remote_blocks:
854 + logger.error(
855 + "req %s: local blocks(%d) less than remote blocks(%d)!",
856 + d_req_id,
857 + num_local_blocks,
858 + num_remote_blocks,
859 + )
860 + err_reqs.append(d_req_id)
861 + continue
862 if num_local_blocks > num_remote_blocks:
863 local_block_ids = local_block_ids[-num_remote_blocks:]
864
@@ -643,12 +883,12 @@ class MooncakeConnectorWorker:
883
884 logger.debug(
885 "Sending kv_caches for request %s (%d blocks) to %s",
646 - req_id,
886 + d_req_id,
887 num_remote_blocks,
888 remote_session,
889 )
890
651 - return src_ptrs, dst_ptrs, lengths
891 + return src_ptrs, dst_ptrs, lengths, err_reqs
892
893 def _send_blocks(
894 self,
@@ -722,15 +962,12 @@ class MooncakeConnectorWorker:
962 )
963
964 # No need to launch server for D node.
725 - if self.kv_role == "kv_consumer":
965 + if self.is_kv_consumer:
966 return
967
968 ready_event = threading.Event()
969 asyncio.run_coroutine_threadsafe(
730 - self._mooncake_sender_listener(
731 - ready_event, self.side_channel_port, self.tp_rank
732 - ),
733 - self.sender_loop,
970 + self._mooncake_sender_listener(ready_event), self.sender_loop
971 )
972 ready_event.wait() # Wait for listener ZMQ socket to be ready.
973
@@ -745,21 +982,25 @@ class MooncakeConnectorWorker:
982
983 # Handle timeout to avoid stranding blocks on remote.
984 now = time.perf_counter()
748 - expired_reqs = [
749 - req_id
750 - for req_id, send_meta in self.reqs_need_send.items()
751 - if send_meta.expire_time < now
752 - ]
753 - for req_id in expired_reqs:
754 - logger.warning(
755 - "Request %s timed out after %d seconds without "
756 - "being sent. Freeing its blocks on the producer side.",
757 - req_id,
758 - envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
759 - )
760 - del self.reqs_need_send[req_id]
761 - if expired_reqs:
762 - finished_sending_reqs.update(expired_reqs)
985 +
986 + expired_transfer_id = []
987 + for transfer_id, send_meta in self.reqs_need_send.items():
988 + if (
989 + send_meta.p_req_id
990 + and send_meta.expire_time < now
991 + and send_meta.sending == 0
992 + ):
993 + logger.warning(
994 + "Request %s timed out after %d seconds without "
995 + "being sent. Freeing its blocks on the producer side.",
996 + send_meta.p_req_id,
997 + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
998 + )
999 + finished_sending_reqs.add(send_meta.p_req_id)
1000 + expired_transfer_id.append(transfer_id)
1001 +
1002 + for transfer_id in expired_transfer_id:
1003 + del self.reqs_need_send[transfer_id]
1004
1005 return finished_sending_reqs
1006
@@ -771,12 +1012,12 @@ class MooncakeConnectorWorker:
1012 """
1013 recv_fut = None
1014 send_fut = None
774 - if self.kv_role != "kv_producer":
1015 + if not self.is_kv_producer:
1016 recv_fut = asyncio.run_coroutine_threadsafe(
1017 self.fetch_finished_recving_reqs(), self.receiver_loop
1018 )
1019
779 - if self.kv_role != "kv_consumer":
1020 + if not self.is_kv_consumer:
1021 send_fut = asyncio.run_coroutine_threadsafe(
1022 self.fetch_finished_sending_reqs(), self.sender_loop
1023 )
@@ -795,69 +1036,174 @@ class MooncakeConnectorWorker:
1036
1037 return finished_sending_reqs or None, finished_recving_reqs or None
1038
798 - async def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]):
799 - req_ids, block_ids = map(list, zip(*req_blocks))
800 - metadata = MooncakeAgentMetadata(
1039 + async def receive_kv_from_single_worker(
1040 + self,
1041 + worker_addr: str,
1042 + pull_metas: dict[ReqId, PullReqMeta],
1043 + ):
1044 + req_ids = set(pull_metas)
1045 + metadata = MooncakeXferMetadata(
1046 remote_hostname=self.hostname,
1047 remote_port=self.rpc_port,
803 - request_ids=req_ids,
1048 + remote_tp_size=self.tp_size,
1049 + remote_tp_rank=self.tp_rank,
1050 + req_blocks={
1051 + req_id: (pull_meta.transfer_id, pull_meta.local_block_ids)
1052 + for req_id, pull_meta in pull_metas.items()
1053 + },
1054 kv_caches_base_addr=self.kv_caches_base_addr,
805 - block_ids=block_ids,
1055 )
1056
1057 encoded_data = self._encoder.encode(metadata)
1058 logger.debug(
810 - "Size of encoded MooncakeAgentMetadata: %d bytes", len(encoded_data)
1059 + "Size of encoded MooncakeXferMetadata: %d bytes", len(encoded_data)
1060 + )
1061 + logger.debug(
1062 + "Sending kv transfer request for %s on path: %s", req_ids, worker_addr
1063 )
812 - logger.debug("Sending kv transfer request for %s on path: %s", req_ids, path)
1064
1065 # Send query for the request.
815 - sock: zmq.asyncio.Socket = make_zmq_socket(
816 - self.async_zmq_ctx, path, zmq.REQ, bind=False, linger=0
817 - )
818 - sock.setsockopt(zmq.RCVTIMEO, 60000)
1066 try:
820 - await sock.send(encoded_data)
821 - ret_msg = await sock.recv()
822 - if ret_msg != TRANS_DONE:
823 - logger.error(
824 - "Error happens during tranfering kvcache for %s, see logs in prefiller.", # noqa: E501
825 - req_ids,
1067 + with make_zmq_socket(
1068 + self.async_zmq_ctx, worker_addr, zmq.DEALER, bind=False, linger=0
1069 + ) as sock:
1070 + # If something goes wrong, let P wait timeout first (in asyncio.wait()).
1071 + sock.setsockopt(
1072 + zmq.RCVTIMEO, (envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT + 60) * 1000
1073 )
827 - return
1074 + await sock.send(encoded_data)
1075 + while True:
1076 + ret_msg = await sock.recv()
1077 + response = self._xfer_resp_decoder.decode(ret_msg)
1078 + if response.status == MooncakeXferResponseStatus.ERROR:
1079 + logger.error(
1080 + "Error happens during tranfering kvcache for %s: %s",
1081 + req_ids,
1082 + response.err_msg,
1083 + )
1084 + return
1085 + self.process_pulling_result(response, pull_metas)
1086 + if response.status == MooncakeXferResponseStatus.FINISH:
1087 + break
1088 except zmq.ContextTerminated:
1089 logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
1090 except Exception as e:
831 - logger.error("MooncakeAgentMetadata transfer failed for %s: %s", req_ids, e)
1091 + logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
1092 return
833 - finally:
834 - sock.close()
1093
836 - self.finished_recving_reqs.update(req_ids)
1094 + def process_pulling_result(
1095 + self,
1096 + response: MooncakeXferResponse,
1097 + pull_metas: dict[ReqId, PullReqMeta],
1098 + ):
1099 + ok_reqs: list[ReqId] = response.ok_reqs or []
1100 +
1101 + for req_id in ok_reqs:
1102 + pull_meta = pull_metas[req_id]
1103 + # No race because we are in async loop.
1104 + pull_meta.pull_tasks_count -= 1
1105 + if pull_meta.pull_tasks_count == 0:
1106 + self.finished_recving_reqs.add(pull_meta.d_req_id)
1107 +
1108 + if ok_reqs:
1109 + logger.debug("pulling kv_caches for %s finished", ok_reqs)
1110 +
1111 + if response.err_reqs:
1112 + logger.error(
1113 + "pulling kv_caches for %s failed: %s",
1114 + response.err_reqs,
1115 + response.err_msg,
1116 + )
1117
838 - logger.debug("pulling kv_caches for %s finished", req_ids)
1118 + async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
1119 + url = remote_bootstrap_addr + "/query"
1120 + try:
1121 + async with httpx.AsyncClient() as client:
1122 + response = await client.get(url)
1123 + response.raise_for_status()
1124 + data: dict = response.json()
1125 + for _, dp_entry in data.items():
1126 + remote_engine_id = dp_entry["engine_id"]
1127 + self._remote_agents[remote_engine_id] = {
1128 + int(tp_rank): {
1129 + int(pp_rank): worker_addr
1130 + for pp_rank, worker_addr in tp_entry.items()
1131 + }
1132 + for tp_rank, tp_entry in dp_entry["worker_addr"].items()
1133 + }
1134 + self._tp_size[remote_engine_id] = len(dp_entry["worker_addr"])
1135 + except Exception as e:
1136 + logger.error(
1137 + "Failed to connect to bootstrap server %s: %s",
1138 + remote_bootstrap_addr,
1139 + e,
1140 + )
1141
840 - def group_kv_pull(self, metadata: MooncakeConnectorMetadata):
841 - kv_pulls = defaultdict(list)
842 - for req_id, meta in metadata.reqs_to_recv.items():
843 - logger.debug(
844 - "start_load_kv for request %s from remote engine. "
845 - "Num local_block_ids: %s.",
846 - req_id,
847 - len(meta.local_block_ids),
1142 + # Always notify others regardless of connection success or failure.
1143 + self._pending_bootstrap_querys[remote_bootstrap_addr].set()
1144 + del self._pending_bootstrap_querys[remote_bootstrap_addr]
1145 +
1146 + def receive_kv(
1147 + self,
1148 + remote_engine_id: EngineId,
1149 + pull_metas: dict[ReqId, PullReqMeta],
1150 + ):
1151 + remote_tp_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
1152 + remote_engine_id
1153 + )
1154 + count = len(remote_tp_ranks)
1155 + if count != 1:
1156 + logger.error("Mooncake: Heterogeneous TP is not supported yet.")
1157 + raise NotImplementedError(
1158 + "Mooncake: Heterogeneous TP is not supported yet."
1159 )
849 - path = make_zmq_path(
850 - "tcp", meta.remote_host, meta.remote_port + self.tp_rank
1160 + for pull_meta in pull_metas.values():
1161 + pull_meta.pull_tasks_count = count
1162 + for remote_tp_rank in remote_tp_ranks:
1163 + worker_addr = self._remote_agents[remote_engine_id][remote_tp_rank][0]
1164 + asyncio.create_task(
1165 + self.receive_kv_from_single_worker(worker_addr, pull_metas)
1166 + )
1167 +
1168 + async def handle_new_engine_id(
1169 + self,
1170 + remote_engine_id: EngineId,
1171 + pull_metas: dict[ReqId, PullReqMeta],
1172 + ):
1173 + remote_bootstrap_addr = next(iter(pull_metas.values())).remote_bootstrap_addr
1174 + if remote_bootstrap_addr not in self._pending_bootstrap_querys:
1175 + self._pending_bootstrap_querys[remote_bootstrap_addr] = asyncio.Event()
1176 + await self._connect_to_prefiller_bootstrap(remote_bootstrap_addr)
1177 + else:
1178 + await self._pending_bootstrap_querys[remote_bootstrap_addr].wait()
1179 +
1180 + if remote_engine_id not in self._remote_agents:
1181 + logger.error(
1182 + "Failed to find remote engine_id %s from bootstrap server %s",
1183 + remote_engine_id,
1184 + remote_bootstrap_addr,
1185 )
852 - kv_pulls[path].append((req_id, meta.local_block_ids))
1186 + return
1187 +
1188 + self.receive_kv(remote_engine_id, pull_metas)
1189
854 - return kv_pulls
1190 + async def _start_load_kv(
1191 + self, reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]]
1192 + ):
1193 + for remote_engine_id, pull_metas in reqs_to_recv.items():
1194 + if remote_engine_id not in self._remote_agents:
1195 + asyncio.create_task(
1196 + self.handle_new_engine_id(remote_engine_id, pull_metas)
1197 + )
1198 + else:
1199 + self.receive_kv(remote_engine_id, pull_metas)
1200
1201 async def record_send_reqs(self, metadata: MooncakeConnectorMetadata):
857 - for req_id, block_ids in metadata.reqs_to_send.items():
1202 + for p_req_id, (transfer_id, block_ids) in metadata.reqs_to_send.items():
1203 if block_ids:
1204 # Already gone through request_finished()
860 - send_meta = self.reqs_need_send[req_id]
1205 + send_meta = self.reqs_need_send[transfer_id]
1206 + send_meta.p_req_id = p_req_id
1207 send_meta.local_block_ids = block_ids
1208 send_meta.expire_time = (
1209 time.perf_counter() + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
@@ -866,20 +1212,29 @@ class MooncakeConnectorWorker:
1212 else:
1213 # From update_state_after_alloc(),
1214 # but not reach request_finished() yet
869 - self.reqs_need_send[req_id] = SendBlockMeta(
870 - local_block_ids=[],
871 - ready=asyncio.Event(),
872 - )
1215 + # This may be already created by send_kv_to_decode()
1216 + # when D is sending MooncakeXferMetadata.
1217 + if transfer_id not in self.reqs_need_send:
1218 + self.reqs_need_send[transfer_id] = SendBlockMeta(
1219 + p_req_id=p_req_id,
1220 + transfer_id=transfer_id,
1221 + local_block_ids=[],
1222 + ready=asyncio.Event(),
1223 + )
1224 + for transfer_id in metadata.reqs_not_processed:
1225 + send_meta = self.reqs_need_send.pop(transfer_id)
1226 + if send_meta:
1227 + assert not send_meta.ready.is_set()
1228
1229 def start_load_kv(self, metadata: MooncakeConnectorMetadata):
875 - if self.kv_role != "kv_producer":
876 - kv_pulls = self.group_kv_pull(metadata)
877 - for path, req_blocks in kv_pulls.items():
878 - asyncio.run_coroutine_threadsafe(
879 - self.receive_kv(path, req_blocks), self.receiver_loop
880 - )
1230 + if not self.is_kv_producer and metadata.reqs_to_recv:
1231 + asyncio.run_coroutine_threadsafe(
1232 + self._start_load_kv(metadata.reqs_to_recv), self.receiver_loop
1233 + )
1234
882 - if self.kv_role != "kv_consumer":
1235 + if not self.is_kv_consumer and (
1236 + metadata.reqs_to_send or metadata.reqs_not_processed
1237 + ):
1238 asyncio.run_coroutine_threadsafe(
1239 self.record_send_reqs(metadata), self.sender_loop
1240 )
@@ -914,3 +1269,31 @@ def get_mooncake_side_channel_port(vllm_config: VllmConfig) -> int:
1269 def _async_loop(loop: asyncio.AbstractEventLoop):
1270 asyncio.set_event_loop(loop)
1271 loop.run_forever()
1272 +
1273 +
1274 +def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool:
1275 + assert (parallel_config := vllm_config.parallel_config)
1276 + # In hybrid or external LB mode,
1277 + # each instance should have its own bootstrap server.
1278 + #
1279 + # In internal LB mode,
1280 + # only the real global first rank need to launch the bootstrap server.
1281 + return is_local_first_rank() and (
1282 + parallel_config.local_engines_only or parallel_config.data_parallel_index == 0
1283 + )
1284 +
1285 +
1286 +def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]:
1287 + """
1288 + Returns the address of the Mooncake bootstrap server.
1289 + This is only used by prefillers to register workers.
1290 + Decoders should get addr from kv_transfer_params.
1291 + """
1292 + assert (parallel_config := vllm_config.parallel_config)
1293 + if parallel_config.local_engines_only:
1294 + # In hybrid or external LB mode, connect to local server.
1295 + host = "127.0.0.1"
1296 + else:
1297 + host = parallel_config.data_parallel_master_ip
1298 + port = envs.VLLM_MOONCAKE_BOOTSTRAP_PORT
1299 + return (host, port)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py new
+127
@@ -0,0 +1,127 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import threading
4 +import time
5 +from dataclasses import dataclass
6 +
7 +import uvicorn
8 +from fastapi import FastAPI, HTTPException
9 +from pydantic import BaseModel
10 +
11 +from vllm.config import VllmConfig
12 +from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
13 +from vllm.logger import init_logger
14 +
15 +WorkerAddr = str
16 +
17 +logger = init_logger(__name__)
18 +
19 +
20 +class RegisterWorkerPayload(BaseModel):
21 + engine_id: EngineId
22 + dp_rank: int
23 + tp_rank: int
24 + pp_rank: int
25 + addr: WorkerAddr
26 +
27 +
28 +@dataclass
29 +class EngineEntry:
30 + engine_id: EngineId
31 + # {tp_rank: {pp_rank: worker_addr}}
32 + worker_addr: dict[int, dict[int, WorkerAddr]]
33 +
34 +
35 +class MooncakeBootstrapServer:
36 + """
37 + A centralized server running on the global rank 0 prefiller worker.
38 + Prefiller workers register their connection info (IP, port, ranks) here.
39 + """
40 +
41 + def __init__(self, vllm_config: VllmConfig, host: str, port: int):
42 + self.workers: dict[int, EngineEntry] = {}
43 +
44 + self.host = host
45 + self.port = port
46 + self.app = FastAPI()
47 + self._register_routes()
48 + self.server_thread: threading.Thread | None = None
49 + self.server: uvicorn.Server | None = None
50 +
51 + def __del__(self):
52 + self.shutdown()
53 +
54 + def _register_routes(self):
55 + # All methods are async. No need to use lock to protect data.
56 + self.app.post("/register")(self.register_worker)
57 + self.app.get("/query", response_model=dict[int, EngineEntry])(self.query)
58 +
59 + def start(self):
60 + if self.server_thread:
61 + return
62 +
63 + config = uvicorn.Config(app=self.app, host=self.host, port=self.port)
64 + self.server = uvicorn.Server(config=config)
65 + self.server_thread = threading.Thread(
66 + target=self.server.run, name="mooncake_bootstrap_server", daemon=True
67 + )
68 + self.server_thread.start()
69 + while not self.server.started:
70 + time.sleep(0.1) # Wait for the server to start
71 + logger.info("Mooncake Bootstrap Server started at %s:%d", self.host, self.port)
72 +
73 + def shutdown(self):
74 + if self.server_thread is None or self.server is None or not self.server.started:
75 + return
76 +
77 + self.server.should_exit = True
78 + self.server_thread.join()
79 + logger.info("Mooncake Bootstrap Server stopped.")
80 +
81 + async def register_worker(self, payload: RegisterWorkerPayload):
82 + """Handles registration of a prefiller worker."""
83 + if payload.dp_rank not in self.workers:
84 + self.workers[payload.dp_rank] = EngineEntry(
85 + engine_id=payload.engine_id,
86 + worker_addr={},
87 + )
88 +
89 + dp_entry = self.workers[payload.dp_rank]
90 + if dp_entry.engine_id != payload.engine_id:
91 + raise HTTPException(
92 + status_code=400,
93 + detail=(
94 + f"Engine ID mismatch for dp_rank={payload.dp_rank}: "
95 + f"expected {dp_entry.engine_id}, got {payload.engine_id}"
96 + ),
97 + )
98 + if payload.tp_rank not in dp_entry.worker_addr:
99 + dp_entry.worker_addr[payload.tp_rank] = {}
100 +
101 + tp_entry = dp_entry.worker_addr[payload.tp_rank]
102 + if payload.pp_rank in tp_entry:
103 + raise HTTPException(
104 + status_code=400,
105 + detail=(
106 + f"Worker with dp_rank={payload.dp_rank}, "
107 + f"tp_rank={payload.tp_rank}, pp_rank={payload.pp_rank} "
108 + f"is already registered at "
109 + f"{tp_entry[payload.pp_rank]}, "
110 + f"but still want to register at {payload.addr}"
111 + ),
112 + )
113 +
114 + tp_entry[payload.pp_rank] = payload.addr
115 + logger.debug(
116 + "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, pp_rank=%d at %s",
117 + payload.engine_id,
118 + payload.dp_rank,
119 + payload.tp_rank,
120 + payload.pp_rank,
121 + payload.addr,
122 + )
123 +
124 + return {"status": "ok"}
125 +
126 + async def query(self) -> dict[int, EngineEntry]:
127 + return self.workers