| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # |
| 4 | # Minimal single-node reference proxy for MoRI-IO prefill/decode disaggregation. |
| 5 | # It demonstrates the DP-rank pinning contract the connector expects: pick one |
| 6 | # prefill DP rank per request (flat_interleaved_dp_route) and pin BOTH legs to it |
| 7 | # via the ``X-data-parallel-rank`` header plus ``kv_transfer_params`` |
| 8 | # (``remote_dp_rank`` / ``remote_dp_size`` / ``remote_tp_size``). |
| 9 | # |
| 10 | # Scope: single-node only. It intentionally does NOT supply the cross-pod |
| 11 | # Wide-EP (2P2D) peer info -- ``remote_hosts`` / ``multi_pod_hosts``, |
| 12 | # ``data_parallel_size_local``, ``is_request_leader``. In real multi-pod |
| 13 | # deployments that peer info is produced by a production routing sidecar, e.g.: |
| 14 | # * llm-d-router (https://github.com/llm-d/llm-d-router) |
| 15 | # * vLLM's own vllm-router (examples in the vllm-project org) |
| 16 | # * or any NIXL-shaped router that sets the same kv_transfer_params contract. |
| 17 | # Use this file as a protocol reference, not a production router. |
| 18 | import argparse |
| 19 | import asyncio |
| 20 | import copy |
| 21 | import logging |
| 22 | import os |
| 23 | import socket |
| 24 | import threading |
| 25 | import uuid |
| 26 | from urllib.parse import urlparse |
| 27 | |
| 28 | import aiohttp |
| 29 | import msgpack |
| 30 | import zmq |
| 31 | from quart import Quart, Request, make_response, request |
| 32 | |
| 33 | from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( |
| 34 | MoRIIOConstants, |
| 35 | ) |
| 36 | |
| 37 | logger = logging.getLogger(__name__) |
| 38 | logger.setLevel(logging.DEBUG) |
| 39 | prefill_instances: list[dict] = [] |
| 40 | decode_instances: list[dict] = [] |
| 41 | request_nums = 0 |
| 42 | app = Quart(__name__) |
| 43 | |
| 44 | |
| 45 | TRANSFER_TYPE = None |
| 46 | |
| 47 | |
| 48 | _list_lock = threading.RLock() |
| 49 | |
| 50 | |
| 51 | def _listen_for_register(hostname, port): |
| 52 | context = zmq.Context() |
| 53 | router_socket = context.socket(zmq.ROUTER) |
| 54 | router_socket.bind(f"tcp://{hostname}:{port}") |
| 55 | poller = zmq.Poller() |
| 56 | poller.register(router_socket, zmq.POLLIN) |
| 57 | global prefill_instances |
| 58 | global decode_instances |
| 59 | |
| 60 | while True: |
| 61 | socks = dict(poller.poll()) |
| 62 | if router_socket in socks: |
| 63 | remote_addr, msg = router_socket.recv_multipart() |
| 64 | data = msgpack.loads(msg) |
| 65 | if data.get("type") == "HELLO": |
| 66 | pass |
| 67 | elif data.get("type") in ("P", "D"): |
| 68 | role = data["type"] |
| 69 | required_keys = { |
| 70 | "http_address", |
| 71 | "zmq_address", |
| 72 | "dp_size", |
| 73 | "tp_size", |
| 74 | "transfer_mode", |
| 75 | } |
| 76 | missing = required_keys - data.keys() |
| 77 | if missing: |
| 78 | logger.error( |
| 79 | "Registration message missing required keys %s; skipping", |
| 80 | missing, |
| 81 | ) |
| 82 | continue |
| 83 | # Derive request_address from http_address |
| 84 | # api path suffix is appended at request time |
| 85 | instance = { |
| 86 | "role": role, |
| 87 | "request_address": f"http://{data['http_address']}/v1", |
| 88 | "http_address": data["http_address"], |
| 89 | "zmq_address": data["zmq_address"], |
| 90 | "dp_size": data["dp_size"], |
| 91 | "tp_size": data["tp_size"], |
| 92 | "transfer_mode": data["transfer_mode"], |
| 93 | } |
| 94 | # zmq_address format: "host:IP,handshake:PORT,notify:PORT" |
| 95 | # Stored verbatim; embedded into the request_id by handle_request. |
| 96 | |
| 97 | global TRANSFER_TYPE |
| 98 | transfer_mode = instance["transfer_mode"] |
| 99 | target_list = prefill_instances if role == "P" else decode_instances |
| 100 | with _list_lock: |
| 101 | if TRANSFER_TYPE is None: |
| 102 | TRANSFER_TYPE = transfer_mode |
| 103 | logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE) |
| 104 | elif transfer_mode != TRANSFER_TYPE: |
| 105 | logger.error( |
| 106 | "Mismatched transfer mode: expected %s, got %s;" |
| 107 | " skipping registration of %s", |
| 108 | TRANSFER_TYPE, |
| 109 | transfer_mode, |
| 110 | data["http_address"], |
| 111 | ) |
| 112 | continue |
| 113 | existing_idx = next( |
| 114 | ( |
| 115 | idx |
| 116 | for idx, i in enumerate(target_list) |
| 117 | if i.get("http_address") == data["http_address"] |
| 118 | ), |
| 119 | None, |
| 120 | ) |
| 121 | if existing_idx is not None: |
| 122 | target_list[existing_idx] = instance |
| 123 | logger.info( |
| 124 | "Updated existing %s instance: %s", |
| 125 | "Prefill" if role == "P" else "Decode", |
| 126 | instance, |
| 127 | ) |
| 128 | else: |
| 129 | target_list.append(instance) |
| 130 | logger.info( |
| 131 | "Registered %s instance: %s", |
| 132 | "Prefill" if role == "P" else "Decode", |
| 133 | instance, |
| 134 | ) |
| 135 | else: |
| 136 | logger.warning( |
| 137 | "Received message with unrecognized type %r; ignoring", |
| 138 | data.get("type"), |
| 139 | ) |
| 140 | |
| 141 | |
| 142 | def start_service_discovery(hostname, port): |
| 143 | if not hostname: |
| 144 | hostname = socket.gethostname() |
| 145 | if port == 0: |
| 146 | raise ValueError("Port cannot be 0") |
| 147 | |
| 148 | _listener_thread = threading.Thread( |
| 149 | target=_listen_for_register, args=(hostname, port), daemon=True |
| 150 | ) |
| 151 | _listener_thread.start() |
| 152 | return _listener_thread |
| 153 | |
| 154 | |
| 155 | async def send_request_to_prefill( |
| 156 | endpoint, req_data, request_id, selected_prefill_dp_rank |
| 157 | ): |
| 158 | req_data_copy = req_data |
| 159 | |
| 160 | req_data_copy["kv_transfer_params"].update( |
| 161 | { |
| 162 | "do_remote_decode": True, |
| 163 | "do_remote_prefill": False, |
| 164 | "remote_engine_id": None, |
| 165 | "remote_block_ids": None, |
| 166 | } |
| 167 | ) |
| 168 | req_data_copy["stream"] = False |
| 169 | req_data_copy["max_tokens"] = 1 |
| 170 | if "max_completion_tokens" in req_data_copy: |
| 171 | req_data_copy["max_completion_tokens"] = 1 |
| 172 | if "stream_options" in req_data_copy: |
| 173 | del req_data_copy["stream_options"] |
| 174 | async with aiohttp.ClientSession( |
| 175 | timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000) |
| 176 | ) as session: |
| 177 | headers = { |
| 178 | "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", |
| 179 | "X-Request-Id": request_id, |
| 180 | } |
| 181 | if selected_prefill_dp_rank is not None: |
| 182 | headers["X-data-parallel-rank"] = str(selected_prefill_dp_rank) |
| 183 | async with session.post( |
| 184 | url=endpoint, json=req_data_copy, headers=headers |
| 185 | ) as response: |
| 186 | if response.status == 200: |
| 187 | return await response.json() |
| 188 | |
| 189 | else: |
| 190 | error_message = ( |
| 191 | f"send_request_to_prefill response ={response}," |
| 192 | f"reason={response.reason}, status={response.status}," |
| 193 | f"method={response.method}, url={response.url}," |
| 194 | f"real_url={response.real_url}" |
| 195 | ) |
| 196 | raise RuntimeError(error_message) |
| 197 | |
| 198 | |
| 199 | async def start_decode_request(endpoint, req_data, request_id, selected_dp_rank=None): |
| 200 | session = aiohttp.ClientSession( |
| 201 | timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000) |
| 202 | ) |
| 203 | headers = { |
| 204 | "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", |
| 205 | "X-Request-Id": request_id, |
| 206 | } |
| 207 | # Pin the decode leg to the same DP rank as its prefill leg so both land on |
| 208 | # the same rank (vLLM reads this header directly; no engine-side routing). |
| 209 | if selected_dp_rank is not None: |
| 210 | headers["X-data-parallel-rank"] = str(selected_dp_rank) |
| 211 | response = await session.post(url=endpoint, json=req_data, headers=headers) |
| 212 | return session, response |
| 213 | |
| 214 | |
| 215 | async def stream_decode_response(session, response, request_id): |
| 216 | try: |
| 217 | if response.status == 200: |
| 218 | async for chunk_bytes in response.content.iter_chunked(1024): |
| 219 | yield chunk_bytes |
| 220 | else: |
| 221 | error_message = ( |
| 222 | f"stream_decode_response response ={response}," |
| 223 | f"reason={response.reason}, status={response.status}," |
| 224 | f"method={response.method}, url={response.url}," |
| 225 | f"real_url={response.real_url}" |
| 226 | ) |
| 227 | raise RuntimeError(error_message) |
| 228 | finally: |
| 229 | await session.close() |
| 230 | |
| 231 | |
| 232 | def flat_interleaved_dp_route(request_number, instances): |
| 233 | """Flat round-robin over the full (instance, dp_rank) slot space. |
| 234 | |
| 235 | ONE counter over (n_instances * dp_size) slots, so instance-selection and |
| 236 | DP-rank-selection are derived from the SAME index and can never alias. The |
| 237 | previous scheme computed instance = req % n and rank = req % dp from the |
| 238 | same counter with n | dp, which locked each instance to a stride-n subset |
| 239 | of its ranks (e.g. 2 prefill instances -> 4 of 8 ranks each -> half the |
| 240 | GPUs never receive a request, so the deployment falsely appears not to |
| 241 | scale). |
| 242 | |
| 243 | Interleaved order — inst0_r0, inst1_r0, inst0_r1, inst1_r1, ... — so |
| 244 | consecutive requests alternate instances AND every rank gets walked. |
| 245 | |
| 246 | Assumes homogeneous dp_size across a role's instances (true for the |
| 247 | DP<->DP and DP<->TP deployments this proxy targets). Returns |
| 248 | (instance_index, dp_rank); dp_rank is None when dp_size == 1 (e.g. a TP |
| 249 | decode), which avoids forwarding an out-of-range data-parallel rank. |
| 250 | """ |
| 251 | n = len(instances) |
| 252 | dp = instances[0]["dp_size"] |
| 253 | slot = (request_number - 1) % (n * dp) |
| 254 | inst_idx = slot % n |
| 255 | dp_rank = (slot // n) if dp > 1 else None |
| 256 | return inst_idx, dp_rank |
| 257 | |
| 258 | |
| 259 | @app.route("/health", methods=["GET"]) |
| 260 | async def health(): |
| 261 | # Benchmark harnesses and load balancers often probe the proxy URL and |
| 262 | # treat a non-200 response as a dead service. Backends expose /health, but |
| 263 | # this proxy used to lack one; report the proxy process itself as healthy. |
| 264 | return ("ok", 200) |
| 265 | |
| 266 | |
| 267 | @app.route("/v1/completions", methods=["POST"]) |
| 268 | async def handle_completions_request(): |
| 269 | return await handle_request("/completions", request) |
| 270 | |
| 271 | |
| 272 | @app.route("/v1/chat/completions", methods=["POST"]) |
| 273 | async def handle_chat_completions_request(): |
| 274 | return await handle_request("/chat/completions", request) |
| 275 | |
| 276 | |
| 277 | async def handle_request(api: str, request: Request): |
| 278 | try: |
| 279 | with _list_lock: |
| 280 | global request_nums |
| 281 | request_nums += 1 |
| 282 | |
| 283 | req_data = await request.get_json() |
| 284 | |
| 285 | prefill_instance_endpoint = None |
| 286 | decode_instance_endpoint = None |
| 287 | error_msg = ( |
| 288 | "Service Unavailable: No prefill or decode instances are registered." |
| 289 | ) |
| 290 | if not prefill_instances or not decode_instances: |
| 291 | return await make_response( |
| 292 | ( |
| 293 | error_msg, |
| 294 | 503, |
| 295 | ) |
| 296 | ) |
| 297 | # Flat interleaved round-robin (see flat_interleaved_dp_route): ONE |
| 298 | # counter over the full (instance, dp_rank) slot space per role, so |
| 299 | # instance-selection and DP-rank-selection derive from the same index |
| 300 | # and can never alias. The old scheme keyed both on request_nums with |
| 301 | # n_instances | dp_size, stranding half the ranks (e.g. in 2P_DP8EP). |
| 302 | pid, selected_prefill_dp_rank = flat_interleaved_dp_route( |
| 303 | request_nums, prefill_instances |
| 304 | ) |
| 305 | # Decode instance selection uses the same interleaved walk; in READ |
| 306 | # mode the decode reads KV from selected_prefill_dp_rank, so the |
| 307 | # decode's own dp_rank is not forwarded here. |
| 308 | did, _ = flat_interleaved_dp_route(request_nums, decode_instances) |
| 309 | prefill_instance_endpoint = prefill_instances[pid] |
| 310 | decode_instance_endpoint = decode_instances[did] |
| 311 | |
| 312 | # Embed both zmq_addresses in the request_id so the connector can parse |
| 313 | # the peer's host/ports from it, similar to P2P-NCCL |
| 314 | uid = str(uuid.uuid4()).replace("-", "") |
| 315 | request_id = ( |
| 316 | f"___prefill_addr_{prefill_instance_endpoint['zmq_address']}" |
| 317 | f"___decode_addr_{decode_instance_endpoint['zmq_address']}" |
| 318 | f"_{uid}" |
| 319 | ) |
| 320 | |
| 321 | transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}" |
| 322 | |
| 323 | req_data_to_prefill = copy.deepcopy(req_data) |
| 324 | req_data_to_prefill["kv_transfer_params"] = {} |
| 325 | req_data["kv_transfer_params"] = {} |
| 326 | req_data_to_prefill["kv_transfer_params"]["remote_dp_size"] = ( |
| 327 | decode_instance_endpoint["dp_size"] |
| 328 | ) |
| 329 | req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = ( |
| 330 | decode_instance_endpoint["tp_size"] |
| 331 | ) |
| 332 | req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id |
| 333 | |
| 334 | prefill_request_url = prefill_instance_endpoint["request_address"] + api |
| 335 | send_prefill_task = asyncio.create_task( |
| 336 | send_request_to_prefill( |
| 337 | prefill_request_url, |
| 338 | req_data_to_prefill, |
| 339 | request_id, |
| 340 | selected_prefill_dp_rank, |
| 341 | ) |
| 342 | ) |
| 343 | |
| 344 | # max_completion_tokens takes precedence when present (that's the limit |
| 345 | # the backend enforces); fall back to max_tokens. If neither is set |
| 346 | # (e.g. benchmark clients that omit it), leave the request unchanged. |
| 347 | if "max_completion_tokens" in req_data: |
| 348 | req_data["max_completion_tokens"] -= 1 |
| 349 | elif "max_tokens" in req_data: |
| 350 | req_data["max_tokens"] -= 1 |
| 351 | |
| 352 | req_data["kv_transfer_params"] = { |
| 353 | "do_remote_decode": False, |
| 354 | "do_remote_prefill": True, |
| 355 | "remote_engine_id": None, |
| 356 | "remote_block_ids": None, |
| 357 | "transfer_id": transfer_id, |
| 358 | } |
| 359 | if TRANSFER_TYPE == "READ": |
| 360 | # In read mode, prefill and decode are executed serially. |
| 361 | prefill_response = await send_prefill_task |
| 362 | prefill_kv = prefill_response["kv_transfer_params"] |
| 363 | req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[ |
| 364 | "remote_engine_id" |
| 365 | ] |
| 366 | req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[ |
| 367 | "remote_block_ids" |
| 368 | ] |
| 369 | req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] |
| 370 | if prefill_kv.get("tp_size") is not None: |
| 371 | req_data["kv_transfer_params"]["tp_size"] = prefill_kv["tp_size"] |
| 372 | |
| 373 | req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[ |
| 374 | "dp_size" |
| 375 | ] |
| 376 | req_data["kv_transfer_params"]["remote_tp_size"] = prefill_instance_endpoint[ |
| 377 | "tp_size" |
| 378 | ] |
| 379 | |
| 380 | if selected_prefill_dp_rank is not None: |
| 381 | req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank |
| 382 | |
| 383 | decode_request_url = decode_instance_endpoint["request_address"] + api |
| 384 | decode_request_task = asyncio.create_task( |
| 385 | start_decode_request( |
| 386 | decode_request_url, req_data, request_id, selected_prefill_dp_rank |
| 387 | ) |
| 388 | ) |
| 389 | |
| 390 | session, decode_response = await decode_request_task |
| 391 | stream_generator = stream_decode_response(session, decode_response, request_id) |
| 392 | response = await make_response(stream_generator) |
| 393 | response.headers["Content-Type"] = decode_response.headers.get( |
| 394 | "Content-Type", "application/json" |
| 395 | ) |
| 396 | return response |
| 397 | except Exception as e: |
| 398 | logger.exception("An error occurred while handling the request: %s", e) |
| 399 | return await make_response( |
| 400 | ( |
| 401 | f"Internal Server Error: {e!s}", |
| 402 | 500, |
| 403 | ) |
| 404 | ) |
| 405 | |
| 406 | |
| 407 | async def send_profile_cmd(req_data: dict, profiler_cmd: str): |
| 408 | assert profiler_cmd in {"start", "stop"} |
| 409 | |
| 410 | with _list_lock: |
| 411 | p_instances = list(prefill_instances) |
| 412 | d_instances = list(decode_instances) |
| 413 | |
| 414 | if not p_instances and not d_instances: |
| 415 | raise RuntimeError( |
| 416 | "Service Unavailable: No prefill or decode instances are registered." |
| 417 | ) |
| 418 | |
| 419 | headers = { |
| 420 | "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", |
| 421 | } |
| 422 | |
| 423 | tasks = [] |
| 424 | |
| 425 | async with aiohttp.ClientSession( |
| 426 | timeout=aiohttp.ClientTimeout(total=60) |
| 427 | ) as session: |
| 428 | for instances in (p_instances, d_instances): |
| 429 | for inst in instances: |
| 430 | _p = urlparse(inst["request_address"]) |
| 431 | url = f"http://{_p.hostname}:{_p.port}/{profiler_cmd}_profile" |
| 432 | |
| 433 | tasks.append( |
| 434 | session.post( |
| 435 | url, |
| 436 | json=req_data, |
| 437 | headers=headers, |
| 438 | ) |
| 439 | ) |
| 440 | |
| 441 | responses = await asyncio.gather(*tasks, return_exceptions=True) |
| 442 | |
| 443 | for r in responses: |
| 444 | if isinstance(r, Exception): |
| 445 | raise r |
| 446 | if r.status >= 400: |
| 447 | msg = await r.text() |
| 448 | raise RuntimeError(f"{profiler_cmd}_profile failed: {r.status}, {msg}") |
| 449 | |
| 450 | return await responses[0].json() |
| 451 | |
| 452 | |
| 453 | @app.post("/start_profile") |
| 454 | async def start_profile(): |
| 455 | try: |
| 456 | req_data = await request.get_json() |
| 457 | return await send_profile_cmd(req_data, "start") |
| 458 | except Exception as e: |
| 459 | logger.exception("start_profile failed: %s", e) |
| 460 | return await make_response((str(e), 500)) |
| 461 | |
| 462 | |
| 463 | @app.post("/stop_profile") |
| 464 | async def stop_profile(): |
| 465 | try: |
| 466 | req_data = await request.get_json() |
| 467 | return await send_profile_cmd(req_data, "stop") |
| 468 | except Exception as e: |
| 469 | logger.exception("stop_profile failed: %s", e) |
| 470 | return await make_response((str(e), 500)) |
| 471 | |
| 472 | |
| 473 | if __name__ == "__main__": |
| 474 | parser = argparse.ArgumentParser() |
| 475 | parser.add_argument("--port", type=int, default=10001) |
| 476 | args = parser.parse_args() |
| 477 | |
| 478 | t = start_service_discovery("0.0.0.0", 36367) |
| 479 | # High-concurrency hardening. Quart's app.run() uses a shallow listen |
| 480 | # backlog (100) and, with app.debug=True, adds per-request overhead that |
| 481 | # starves the single accept loop. Under a burst of ~512 simultaneous client |
| 482 | # connections the backlog overflows and the kernel RSTs the excess, so |
| 483 | # clients see "ClientOSError: [Errno 104] Connection reset by peer" before |
| 484 | # any response (~16% request loss at c=512). Serve via hypercorn with debug |
| 485 | # OFF and a deep backlog so the burst QUEUES (higher TTFT) instead of being |
| 486 | # reset -> 100% request success. |
| 487 | app.debug = False |
| 488 | app.config["BODY_TIMEOUT"] = 360000 |
| 489 | app.config["RESPONSE_TIMEOUT"] = 360000 |
| 490 | |
| 491 | import asyncio |
| 492 | import os |
| 493 | |
| 494 | from hypercorn.asyncio import serve as _hypercorn_serve |
| 495 | from hypercorn.config import Config as _HypercornConfig |
| 496 | |
| 497 | _hcfg = _HypercornConfig() |
| 498 | _hcfg.bind = [f"0.0.0.0:{args.port}"] |
| 499 | # Deep listen backlog so a wide connection burst queues, not RSTs. NOTE: |
| 500 | # effective backlog is capped by the host's net.core.somaxconn (proxy runs |
| 501 | # --network host); kernel 6.x defaults to 4096. Override via |
| 502 | # PROXY_LISTEN_BACKLOG. |
| 503 | _hcfg.backlog = int(os.environ.get("PROXY_LISTEN_BACKLOG", "4096")) |
| 504 | # Long-lived SSE streams (8k1k decode ~5 min): never reap on keepalive. |
| 505 | _hcfg.keep_alive_timeout = 360000.0 |
| 506 | |
| 507 | asyncio.run(_hypercorn_serve(app, _hcfg)) |
| 508 | t.join() |