| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Push-mode disaggregated prefilling proxy demo. |
| 4 | |
| 5 | Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is |
| 6 | the same; the difference is in how P and D coordinate the KV transfer: |
| 7 | |
| 8 | * Pull mode: proxy forwards P's ``kv_transfer_params`` (including |
| 9 | ``remote_block_ids``) to D, and D pulls KV from P via NIXL READ. |
| 10 | * Push mode: proxy hands D **only** P's coordinates |
| 11 | (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``, |
| 12 | ``pp_size``) and the shared ``remote_request_id``. D registers its locally |
| 13 | allocated blocks with P over a NIXL notification; P then pushes the KV to D via |
| 14 | NIXL WRITE. |
| 15 | |
| 16 | Launch multiple vLLM instances configured with ``NixlPushConnector`` and |
| 17 | matching ``engine_id`` / ``side_channel_port``, then start this proxy: |
| 18 | |
| 19 | python3 examples/disaggregated/disaggregated_serving/\ |
| 20 | disagg_proxy_pushconnector_demo.py \ |
| 21 | --model $model_name \ |
| 22 | --prefill localhost:8100 \ |
| 23 | --decode localhost:8200 \ |
| 24 | --prefill-engine-id prefill-engine-001 \ |
| 25 | --prefill-kv-host 10.0.0.1 \ |
| 26 | --prefill-side-channel-port 5600 \ |
| 27 | --prefill-tp-size 1 \ |
| 28 | --prefill-pp-size 1 \ |
| 29 | --port 8000 |
| 30 | """ |
| 31 | |
| 32 | import argparse |
| 33 | import asyncio |
| 34 | import contextlib |
| 35 | import ipaddress |
| 36 | import itertools |
| 37 | import json |
| 38 | import logging |
| 39 | import os |
| 40 | import sys |
| 41 | import uuid |
| 42 | from abc import ABC, abstractmethod |
| 43 | from collections.abc import Callable |
| 44 | |
| 45 | import aiohttp |
| 46 | import uvicorn |
| 47 | from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request |
| 48 | from fastapi.responses import JSONResponse, StreamingResponse |
| 49 | |
| 50 | AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) |
| 51 | logger = logging.getLogger() |
| 52 | logging.basicConfig(level=logging.INFO) |
| 53 | |
| 54 | |
| 55 | class SchedulingPolicy(ABC): |
| 56 | @abstractmethod |
| 57 | def schedule(self, cycler: itertools.cycle): |
| 58 | raise NotImplementedError("Scheduling Proxy is not set.") |
| 59 | |
| 60 | |
| 61 | class RoundRobinSchedulingPolicy(SchedulingPolicy): |
| 62 | def schedule(self, cycler: itertools.cycle) -> str: |
| 63 | return next(cycler) |
| 64 | |
| 65 | |
| 66 | class PushProxy: |
| 67 | """Push-mode proxy. |
| 68 | |
| 69 | The structure mirrors the pull-mode ``Proxy`` in |
| 70 | ``disagg_proxy_demo.py``: an APIRouter with ``/v1/completions``, |
| 71 | ``/v1/chat/completions``, ``/status`` and ``/instances/add``, plus |
| 72 | round-robin scheduling across multiple P / D instances. |
| 73 | |
| 74 | Push-specific differences are confined to the request-handling |
| 75 | methods (``create_completion`` / ``create_chat_completion``): |
| 76 | |
| 77 | * D's ``kv_transfer_params`` is built from CLI-provided P |
| 78 | coordinates instead of being derived from P's response. |
| 79 | * P and D requests are issued concurrently — D registers blocks and |
| 80 | waits while P prefills and pushes. |
| 81 | """ |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | prefill_instances: list[str], |
| 86 | decode_instances: list[str], |
| 87 | model: str, |
| 88 | scheduling_policy: SchedulingPolicy, |
| 89 | prefill_engine_id: str, |
| 90 | prefill_kv_host: str, |
| 91 | prefill_side_channel_port: int, |
| 92 | prefill_tp_size: int, |
| 93 | prefill_pp_size: int, |
| 94 | custom_create_completion: Callable[[Request], StreamingResponse] | None = None, |
| 95 | custom_create_chat_completion: Callable[[Request], StreamingResponse] |
| 96 | | None = None, |
| 97 | ): |
| 98 | self.prefill_instances = prefill_instances |
| 99 | self.decode_instances = decode_instances |
| 100 | self.prefill_cycler = itertools.cycle(prefill_instances) |
| 101 | self.decode_cycler = itertools.cycle(decode_instances) |
| 102 | self.model = model |
| 103 | self.scheduling_policy = scheduling_policy |
| 104 | |
| 105 | # Push-mode metadata: D needs P's coordinates up-front. Pull mode |
| 106 | # learns these from P's response; push mode uses CLI args because |
| 107 | # D issues its registration before P responds. |
| 108 | self.push_metadata = { |
| 109 | "do_remote_decode": False, |
| 110 | "do_remote_prefill": True, |
| 111 | "remote_engine_id": prefill_engine_id, |
| 112 | "remote_host": prefill_kv_host, |
| 113 | "remote_port": prefill_side_channel_port, |
| 114 | "tp_size": prefill_tp_size, |
| 115 | "pp_size": prefill_pp_size, |
| 116 | } |
| 117 | |
| 118 | self.custom_create_completion = custom_create_completion |
| 119 | self.custom_create_chat_completion = custom_create_chat_completion |
| 120 | self.router = APIRouter() |
| 121 | self.setup_routes() |
| 122 | |
| 123 | # ── routes ──────────────────────────────────────────────────────── # |
| 124 | |
| 125 | def setup_routes(self): |
| 126 | self.router.post( |
| 127 | "/v1/completions", dependencies=[Depends(self.validate_json_request)] |
| 128 | )( |
| 129 | self.custom_create_completion |
| 130 | if self.custom_create_completion |
| 131 | else self.create_completion |
| 132 | ) |
| 133 | self.router.post( |
| 134 | "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)] |
| 135 | )( |
| 136 | self.custom_create_chat_completion |
| 137 | if self.custom_create_chat_completion |
| 138 | else self.create_chat_completion |
| 139 | ) |
| 140 | self.router.get("/status", response_class=JSONResponse)(self.get_status) |
| 141 | |
| 142 | async def validate_json_request(self, raw_request: Request): |
| 143 | content_type = raw_request.headers.get("content-type", "").lower() |
| 144 | if content_type != "application/json": |
| 145 | raise HTTPException( |
| 146 | status_code=415, |
| 147 | detail="Unsupported Media Type: Only 'application/json' is allowed", |
| 148 | ) |
| 149 | |
| 150 | # ── HTTP forwarding ─────────────────────────────────────────────── # |
| 151 | |
| 152 | async def forward_request(self, url, data, headers, use_chunked=True): |
| 153 | async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: |
| 154 | try: |
| 155 | async with session.post( |
| 156 | url=url, json=data, headers=headers |
| 157 | ) as response: |
| 158 | if 200 <= response.status < 300 or 400 <= response.status < 500: |
| 159 | if use_chunked: |
| 160 | async for chunk_bytes in response.content.iter_chunked( |
| 161 | 1024 |
| 162 | ): |
| 163 | yield chunk_bytes |
| 164 | else: |
| 165 | yield await response.read() |
| 166 | else: |
| 167 | error_content = await response.text() |
| 168 | with contextlib.suppress(json.JSONDecodeError): |
| 169 | error_content = json.loads(error_content) |
| 170 | logger.error( |
| 171 | "Request failed with status %s: %s", |
| 172 | response.status, |
| 173 | error_content, |
| 174 | ) |
| 175 | raise HTTPException( |
| 176 | status_code=response.status, |
| 177 | detail=f"Request failed with status {response.status}: " |
| 178 | f"{error_content}", |
| 179 | ) |
| 180 | except aiohttp.ClientError as e: |
| 181 | logger.error("ClientError occurred: %s", str(e)) |
| 182 | raise HTTPException( |
| 183 | status_code=502, |
| 184 | detail="Bad Gateway: Error communicating with upstream server.", |
| 185 | ) from e |
| 186 | except Exception as e: |
| 187 | logger.error("Unexpected error: %s", str(e)) |
| 188 | raise HTTPException(status_code=500, detail=str(e)) from e |
| 189 | |
| 190 | def schedule(self, cycler: itertools.cycle) -> str: |
| 191 | return self.scheduling_policy.schedule(cycler) |
| 192 | |
| 193 | async def get_status(self): |
| 194 | return { |
| 195 | "mode": "push", |
| 196 | "prefill_node_count": len(self.prefill_instances), |
| 197 | "decode_node_count": len(self.decode_instances), |
| 198 | "prefill_nodes": self.prefill_instances, |
| 199 | "decode_nodes": self.decode_instances, |
| 200 | "prefill_engine_id": self.push_metadata["remote_engine_id"], |
| 201 | "prefill_kv_host": self.push_metadata["remote_host"], |
| 202 | "prefill_side_channel_port": self.push_metadata["remote_port"], |
| 203 | "prefill_tp_size": self.push_metadata["tp_size"], |
| 204 | "prefill_pp_size": self.push_metadata["pp_size"], |
| 205 | } |
| 206 | |
| 207 | # ── push-mode request handling ──────────────────────────────────── # |
| 208 | |
| 209 | def _build_decode_kv_params(self, request_id: str) -> dict: |
| 210 | """Push-mode kv_transfer_params for D. |
| 211 | |
| 212 | ``remote_block_ids`` is intentionally omitted: D allocates its |
| 213 | own blocks and registers them with P; P determines the |
| 214 | prefill-side block IDs and ships them via the WRITE. |
| 215 | """ |
| 216 | params = self.push_metadata.copy() |
| 217 | params["remote_request_id"] = request_id |
| 218 | return params |
| 219 | |
| 220 | def _common_headers(self, request_id: str) -> dict: |
| 221 | h = {"X-Request-Id": request_id} |
| 222 | api_key = os.environ.get("OPENAI_API_KEY") |
| 223 | if api_key: |
| 224 | h["Authorization"] = f"Bearer {api_key}" |
| 225 | return h |
| 226 | |
| 227 | async def _push_completion(self, raw_request: Request, path: str): |
| 228 | """Shared body for /v1/completions and /v1/chat/completions. |
| 229 | |
| 230 | Push mode fires P and D concurrently: |
| 231 | * P runs a normal prefill (max_tokens=1, do_remote_decode=True). |
| 232 | * D runs the decode (do_remote_prefill=True, no remote_block_ids). |
| 233 | |
| 234 | D blocks waiting for P's WRITE; the response streamed back to the |
| 235 | client is the decode output from D. |
| 236 | """ |
| 237 | request = await raw_request.json() |
| 238 | request_id = str(uuid.uuid4()) |
| 239 | |
| 240 | # Prefill leg (max_tokens=1, signals P to keep KV around for D). |
| 241 | prefill_request = request.copy() |
| 242 | prefill_request["max_tokens"] = 1 |
| 243 | if "max_completion_tokens" in prefill_request: |
| 244 | prefill_request["max_completion_tokens"] = 1 |
| 245 | prefill_request["kv_transfer_params"] = { |
| 246 | "do_remote_decode": True, |
| 247 | "do_remote_prefill": False, |
| 248 | "remote_engine_id": None, |
| 249 | "remote_block_ids": None, |
| 250 | "remote_host": None, |
| 251 | "remote_port": None, |
| 252 | } |
| 253 | |
| 254 | # Decode leg (push mode: no remote_block_ids). |
| 255 | decode_request = request.copy() |
| 256 | decode_request["kv_transfer_params"] = self._build_decode_kv_params(request_id) |
| 257 | |
| 258 | prefill_instance = self.schedule(self.prefill_cycler) |
| 259 | decode_instance = self.schedule(self.decode_cycler) |
| 260 | headers = self._common_headers(request_id) |
| 261 | |
| 262 | async def drain_prefill(): |
| 263 | async for _ in self.forward_request( |
| 264 | f"http://{prefill_instance}{path}", prefill_request, headers |
| 265 | ): |
| 266 | continue |
| 267 | |
| 268 | prefill_task = asyncio.create_task(drain_prefill()) |
| 269 | |
| 270 | async def stream_decode(): |
| 271 | try: |
| 272 | async for chunk in self.forward_request( |
| 273 | f"http://{decode_instance}{path}", decode_request, headers |
| 274 | ): |
| 275 | yield chunk |
| 276 | finally: |
| 277 | await prefill_task |
| 278 | |
| 279 | return StreamingResponse(stream_decode(), media_type="application/json") |
| 280 | |
| 281 | async def create_completion(self, raw_request: Request): |
| 282 | try: |
| 283 | return await self._push_completion(raw_request, "/v1/completions") |
| 284 | except HTTPException: |
| 285 | raise |
| 286 | except Exception: |
| 287 | exc_info = sys.exc_info() |
| 288 | print("Error occurred in disagg push proxy server") |
| 289 | print(exc_info) |
| 290 | raise |
| 291 | |
| 292 | async def create_chat_completion(self, raw_request: Request): |
| 293 | try: |
| 294 | return await self._push_completion(raw_request, "/v1/chat/completions") |
| 295 | except HTTPException: |
| 296 | raise |
| 297 | except Exception: |
| 298 | exc_info = sys.exc_info() |
| 299 | error_messages = [str(e) for e in exc_info if e] |
| 300 | print("Error occurred in disagg push proxy server") |
| 301 | print(error_messages) |
| 302 | return StreamingResponse( |
| 303 | content=iter(error_messages), media_type="text/event-stream" |
| 304 | ) |
| 305 | |
| 306 | |
| 307 | class PushProxyServer: |
| 308 | def __init__( |
| 309 | self, |
| 310 | args: argparse.Namespace, |
| 311 | scheduling_policy: SchedulingPolicy | None = None, |
| 312 | create_completion: Callable[[Request], StreamingResponse] | None = None, |
| 313 | create_chat_completion: Callable[[Request], StreamingResponse] | None = None, |
| 314 | ): |
| 315 | self.validate_parsed_serve_args(args) |
| 316 | self.port = args.port |
| 317 | self.proxy_instance = PushProxy( |
| 318 | prefill_instances=[] if args.prefill is None else args.prefill, |
| 319 | decode_instances=[] if args.decode is None else args.decode, |
| 320 | model=args.model, |
| 321 | scheduling_policy=( |
| 322 | scheduling_policy |
| 323 | if scheduling_policy is not None |
| 324 | else RoundRobinSchedulingPolicy() |
| 325 | ), |
| 326 | prefill_engine_id=args.prefill_engine_id, |
| 327 | prefill_kv_host=args.prefill_kv_host, |
| 328 | prefill_side_channel_port=args.prefill_side_channel_port, |
| 329 | prefill_tp_size=args.prefill_tp_size, |
| 330 | prefill_pp_size=args.prefill_pp_size, |
| 331 | custom_create_completion=create_completion, |
| 332 | custom_create_chat_completion=create_chat_completion, |
| 333 | ) |
| 334 | |
| 335 | def validate_parsed_serve_args(self, args: argparse.Namespace): |
| 336 | if not args.prefill: |
| 337 | raise ValueError("Please specify at least one prefill node.") |
| 338 | if not args.decode: |
| 339 | raise ValueError("Please specify at least one decode node.") |
| 340 | if not args.prefill_engine_id: |
| 341 | raise ValueError( |
| 342 | "--prefill-engine-id is required in push mode (it must match " |
| 343 | "the engine_id passed to the prefill vLLM instance via " |
| 344 | "--kv-transfer-config)." |
| 345 | ) |
| 346 | if not args.prefill_kv_host: |
| 347 | raise ValueError( |
| 348 | "--prefill-kv-host is required in push mode (the IP / host " |
| 349 | "that the prefill vLLM advertises on its NIXL side channel)." |
| 350 | ) |
| 351 | self.validate_instances(args.prefill) |
| 352 | self.validate_instances(args.decode) |
| 353 | |
| 354 | def validate_instances(self, instances: list): |
| 355 | for instance in instances: |
| 356 | if len(instance.split(":")) != 2: |
| 357 | raise ValueError(f"Invalid instance format: {instance}") |
| 358 | host, port = instance.split(":") |
| 359 | try: |
| 360 | if host != "localhost": |
| 361 | ipaddress.ip_address(host) |
| 362 | port = int(port) |
| 363 | if not (0 < port < 65536): |
| 364 | raise ValueError(f"Invalid port number in instance: {instance}") |
| 365 | except Exception as e: |
| 366 | raise ValueError(f"Invalid instance {instance}: {str(e)}") from e |
| 367 | |
| 368 | def run_server(self): |
| 369 | app = FastAPI() |
| 370 | app.include_router(self.proxy_instance.router) |
| 371 | config = uvicorn.Config(app, port=self.port, loop="uvloop") |
| 372 | server = uvicorn.Server(config) |
| 373 | server.run() |
| 374 | |
| 375 | |
| 376 | def parse_args(): |
| 377 | parser = argparse.ArgumentParser("vLLM disaggregated push-mode proxy server.") |
| 378 | parser.add_argument("--model", "-m", type=str, required=True, help="Model name") |
| 379 | |
| 380 | parser.add_argument( |
| 381 | "--prefill", |
| 382 | "-p", |
| 383 | type=str, |
| 384 | nargs="+", |
| 385 | help="List of prefill node URLs (host:port)", |
| 386 | ) |
| 387 | |
| 388 | parser.add_argument( |
| 389 | "--decode", |
| 390 | "-d", |
| 391 | type=str, |
| 392 | nargs="+", |
| 393 | help="List of decode node URLs (host:port)", |
| 394 | ) |
| 395 | |
| 396 | parser.add_argument( |
| 397 | "--port", |
| 398 | type=int, |
| 399 | default=8000, |
| 400 | help="Server port number", |
| 401 | ) |
| 402 | |
| 403 | # Push-mode specific: P's coordinates that D needs in advance. |
| 404 | parser.add_argument( |
| 405 | "--prefill-engine-id", |
| 406 | type=str, |
| 407 | required=True, |
| 408 | help=( |
| 409 | "engine_id of the prefill vLLM instance (must match " |
| 410 | "--kv-transfer-config engine_id on the prefill server)" |
| 411 | ), |
| 412 | ) |
| 413 | parser.add_argument( |
| 414 | "--prefill-kv-host", |
| 415 | type=str, |
| 416 | required=True, |
| 417 | help=( |
| 418 | "IP / host the prefill vLLM advertises on its NIXL side " |
| 419 | "channel (VLLM_NIXL_SIDE_CHANNEL_HOST)" |
| 420 | ), |
| 421 | ) |
| 422 | parser.add_argument( |
| 423 | "--prefill-side-channel-port", |
| 424 | type=int, |
| 425 | default=5600, |
| 426 | help="NIXL side channel port on the prefill node " |
| 427 | "(VLLM_NIXL_SIDE_CHANNEL_PORT, default 5600)", |
| 428 | ) |
| 429 | parser.add_argument( |
| 430 | "--prefill-tp-size", |
| 431 | type=int, |
| 432 | default=1, |
| 433 | help="Tensor parallel size of the prefill vLLM instance", |
| 434 | ) |
| 435 | parser.add_argument( |
| 436 | "--prefill-pp-size", |
| 437 | type=int, |
| 438 | default=1, |
| 439 | help="Pipeline parallel size of the prefill vLLM instance", |
| 440 | ) |
| 441 | return parser.parse_args() |
| 442 | |
| 443 | |
| 444 | if __name__ == "__main__": |
| 445 | args = parse_args() |
| 446 | proxy_server = PushProxyServer(args=args) |
| 447 | proxy_server.run_server() |