| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer. |
| 4 | |
| 5 | This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment, |
| 6 | routing multi-turn chat requests so that each turn reuses KV cache blocks |
| 7 | from the previous turn's Decode node via bidirectional KV transfer. |
| 8 | |
| 9 | Architecture: |
| 10 | Client ──► Proxy ──► Prefill (P) ──► Decode (D) |
| 11 | │ │ │ |
| 12 | │ kv_transfer_params flow: │ |
| 13 | │ D finish ──► proxy caches │ |
| 14 | │ next turn ──► proxy sends │ |
| 15 | │ cached D blocks to P ──► │ |
| 16 | │ P reads D blocks (bidir) │ |
| 17 | │ P sends its blocks to D │ |
| 18 | |
| 19 | Per-request flow: |
| 20 | 1. Client sends chat/completions request to proxy. |
| 21 | 2. Proxy looks up cached D block info from the previous turn |
| 22 | (keyed by conversation_id). |
| 23 | 3. If cache hit, proxy attaches D's block info to the request |
| 24 | so P can read D's KV blocks instead of recomputing. |
| 25 | 4. Proxy sends request to P (max_tokens=1, non-streaming). |
| 26 | 5. P returns kv_transfer_params with its own block info. |
| 27 | 6. Proxy forwards request + P's block info to D (streaming). |
| 28 | 7. D streams the response. The final chunk includes D's |
| 29 | kv_transfer_params, which the proxy caches for the next turn. |
| 30 | 8. Proxy returns D's response to the client. |
| 31 | |
| 32 | Conversation isolation: |
| 33 | Each request must include a ``conversation_id`` field (top-level in |
| 34 | the JSON body) to scope the KV cache across turns. Without it, the |
| 35 | proxy cannot link turns and falls back to no-cache behavior. |
| 36 | |
| 37 | ``conversation_id`` is a non-standard extension to the OpenAI Chat |
| 38 | Completions schema, consumed by this proxy and not forwarded to the |
| 39 | vLLM engine. Strict OpenAI-compatible frontends reject unknown |
| 40 | fields, so clients must opt in only when targeting this proxy. |
| 41 | |
| 42 | Usage: |
| 43 | python disagg_proxy_multiturn.py \\ |
| 44 | --host 0.0.0.0 --port 8000 \\ |
| 45 | --prefiller-host 10.0.0.1 --prefiller-port 8100 \\ |
| 46 | --decoder-host 10.0.0.2 --decoder-port 8200 |
| 47 | |
| 48 | Benchmarking: |
| 49 | Use ``benchmarks/multi_turn/benchmark_serving_multi_turn.py`` with |
| 50 | the ``--send-conversation-id`` flag to inject a per-conversation |
| 51 | ``conversation_id`` into every request so this proxy can key |
| 52 | cross-turn KV cache reuse. The flag is *off by default*: without |
| 53 | it the benchmark sends OpenAI-schema-compliant payloads and every |
| 54 | turn lands as a cache MISS in this proxy. |
| 55 | |
| 56 | Example: |
| 57 | python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\ |
| 58 | --model <MODEL> --served-model-name <NAME> \\ |
| 59 | --url http://<proxy_host>:8000 \\ |
| 60 | --input-file generate_multi_turn.json \\ |
| 61 | --num-clients 2 --max-active-conversations 6 \\ |
| 62 | --send-conversation-id |
| 63 | |
| 64 | See ``docs/features/nixl_connector_usage.md`` for the broader |
| 65 | bidirectional-KV-transfer setup these benchmarks exercise. |
| 66 | |
| 67 | Dependencies: |
| 68 | pip install fastapi uvicorn httpx |
| 69 | |
| 70 | """ |
| 71 | |
| 72 | from __future__ import annotations |
| 73 | |
| 74 | import argparse |
| 75 | import itertools |
| 76 | import json |
| 77 | import logging |
| 78 | import os |
| 79 | import time |
| 80 | import uuid |
| 81 | from contextlib import asynccontextmanager |
| 82 | from dataclasses import dataclass, field |
| 83 | from typing import Any |
| 84 | |
| 85 | import httpx |
| 86 | from fastapi import FastAPI, Request |
| 87 | from fastapi.responses import JSONResponse, StreamingResponse |
| 88 | |
| 89 | # Logging |
| 90 | logging.basicConfig( |
| 91 | format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", |
| 92 | datefmt="%Y-%m-%d %H:%M:%S", |
| 93 | level=logging.INFO, |
| 94 | ) |
| 95 | logger = logging.getLogger("disagg_proxy") |
| 96 | |
| 97 | |
| 98 | # Data structures |
| 99 | @dataclass |
| 100 | class CachedKVEntry: |
| 101 | """KV transfer parameters cached from D's response for one turn.""" |
| 102 | |
| 103 | kv_transfer_params: dict[str, Any] |
| 104 | timestamp: float = field(default_factory=time.time) |
| 105 | |
| 106 | |
| 107 | class ConversationKVCache: |
| 108 | """Per-conversation KV block cache. |
| 109 | |
| 110 | Each conversation is identified by a ``conversation_id`` supplied by |
| 111 | the client. After D finishes a turn, its ``kv_transfer_params`` are |
| 112 | stored here. On the next turn, the proxy retrieves them so P can |
| 113 | read D's blocks via bidirectional KV transfer. |
| 114 | """ |
| 115 | |
| 116 | def __init__(self, ttl_seconds: float = 600.0) -> None: |
| 117 | self._store: dict[str, CachedKVEntry] = {} |
| 118 | self._ttl = ttl_seconds |
| 119 | |
| 120 | def get(self, conversation_id: str) -> dict[str, Any] | None: |
| 121 | """Retrieve and consume cached KV params for a conversation. |
| 122 | |
| 123 | Returns a *copy* of the kv_transfer_params dict, or None. |
| 124 | The entry is removed after retrieval (single-use). |
| 125 | """ |
| 126 | entry = self._store.pop(conversation_id, None) |
| 127 | if entry is None: |
| 128 | return None |
| 129 | age = time.time() - entry.timestamp |
| 130 | if age > self._ttl: |
| 131 | logger.info( |
| 132 | "conv=%s: stale cache entry (age=%.1fs > ttl=%.1fs), discarding", |
| 133 | conversation_id, |
| 134 | age, |
| 135 | self._ttl, |
| 136 | ) |
| 137 | return None |
| 138 | logger.info( |
| 139 | "conv=%s: cache HIT (age=%.1fs)", |
| 140 | conversation_id, |
| 141 | age, |
| 142 | ) |
| 143 | return dict(entry.kv_transfer_params) |
| 144 | |
| 145 | def put(self, conversation_id: str, kv_params: dict[str, Any]) -> None: |
| 146 | """Store D's kv_transfer_params for a conversation.""" |
| 147 | self._store[conversation_id] = CachedKVEntry( |
| 148 | kv_transfer_params=dict(kv_params), # defensive copy |
| 149 | ) |
| 150 | logger.info( |
| 151 | "conv=%s: cached D blocks (remote_request_id=%s, blocks=%d)", |
| 152 | conversation_id, |
| 153 | kv_params.get("remote_request_id", "?"), |
| 154 | len(kv_params.get("remote_block_ids", [[]])[0]) |
| 155 | if kv_params.get("remote_block_ids") |
| 156 | else 0, |
| 157 | ) |
| 158 | |
| 159 | def evict_stale(self) -> int: |
| 160 | """Remove entries older than TTL. Returns count of evicted entries.""" |
| 161 | now = time.time() |
| 162 | stale = [ |
| 163 | cid |
| 164 | for cid, entry in self._store.items() |
| 165 | if now - entry.timestamp > self._ttl |
| 166 | ] |
| 167 | for cid in stale: |
| 168 | del self._store[cid] |
| 169 | return len(stale) |
| 170 | |
| 171 | @property |
| 172 | def size(self) -> int: |
| 173 | return len(self._store) |
| 174 | |
| 175 | |
| 176 | # Global state |
| 177 | kv_cache = ConversationKVCache( |
| 178 | ttl_seconds=450.0 |
| 179 | ) # Must be < VLLM_NIXL_ABORT_REQUEST_TIMEOUT (480s) |
| 180 | |
| 181 | |
| 182 | # Service client helpers |
| 183 | @dataclass |
| 184 | class ServiceClient: |
| 185 | """Wrapper around an httpx.AsyncClient for a P or D instance.""" |
| 186 | |
| 187 | client: httpx.AsyncClient |
| 188 | host: str |
| 189 | port: int |
| 190 | id: int |
| 191 | |
| 192 | |
| 193 | def _make_headers(request_id: str) -> dict[str, str]: |
| 194 | """Build HTTP headers for upstream requests.""" |
| 195 | headers = {"X-Request-Id": request_id} |
| 196 | api_key = os.environ.get("OPENAI_API_KEY") |
| 197 | if api_key: |
| 198 | headers["Authorization"] = f"Bearer {api_key}" |
| 199 | return headers |
| 200 | |
| 201 | |
| 202 | async def _send_to_prefill( |
| 203 | client: ServiceClient, |
| 204 | endpoint: str, |
| 205 | req_data: dict[str, Any], |
| 206 | request_id: str, |
| 207 | ) -> dict[str, Any]: |
| 208 | """Send a non-streaming prefill request (max_tokens=1). |
| 209 | |
| 210 | Returns the JSON response from P, which includes kv_transfer_params. |
| 211 | """ |
| 212 | payload = req_data.copy() |
| 213 | payload["stream"] = False |
| 214 | payload["max_tokens"] = 1 |
| 215 | payload.pop("max_completion_tokens", None) |
| 216 | payload.pop("min_tokens", None) |
| 217 | payload.pop("stream_options", None) |
| 218 | |
| 219 | resp = await client.client.post( |
| 220 | endpoint, |
| 221 | json=payload, |
| 222 | headers=_make_headers(request_id), |
| 223 | ) |
| 224 | resp.raise_for_status() |
| 225 | return resp.json() |
| 226 | |
| 227 | |
| 228 | async def _stream_from_decode( |
| 229 | client: ServiceClient, |
| 230 | endpoint: str, |
| 231 | req_data: dict[str, Any], |
| 232 | request_id: str, |
| 233 | conversation_id: str, |
| 234 | ) -> tuple[str, str | None, dict[str, Any] | None, str, str | None, int | None]: |
| 235 | """Stream response from D, capturing text and kv_transfer_params. |
| 236 | |
| 237 | Returns (collected_text, finish_reason, kv_params, response_id, created). |
| 238 | Also stores kv_params in the conversation cache. |
| 239 | """ |
| 240 | payload = req_data.copy() |
| 241 | payload["stream"] = True |
| 242 | |
| 243 | collected_text = "" |
| 244 | finish_reason: str | None = None |
| 245 | response_id: str | None = None |
| 246 | model_name: str | None = None |
| 247 | created: int | None = None |
| 248 | captured_kv: dict[str, Any] | None = None |
| 249 | |
| 250 | async with client.client.stream( |
| 251 | "POST", |
| 252 | endpoint, |
| 253 | json=payload, |
| 254 | headers=_make_headers(request_id), |
| 255 | ) as resp: |
| 256 | resp.raise_for_status() |
| 257 | async for line in resp.aiter_lines(): |
| 258 | if not line or not line.startswith("data: "): |
| 259 | continue |
| 260 | if line == "data: [DONE]": |
| 261 | break |
| 262 | try: |
| 263 | chunk = json.loads(line[6:]) |
| 264 | except json.JSONDecodeError: |
| 265 | continue |
| 266 | |
| 267 | if response_id is None: |
| 268 | response_id = chunk.get("id") |
| 269 | model_name = chunk.get("model") |
| 270 | created = chunk.get("created") |
| 271 | |
| 272 | for choice in chunk.get("choices", []): |
| 273 | collected_text += choice.get("text", "") |
| 274 | delta = choice.get("delta", {}) |
| 275 | collected_text += delta.get("content", "") |
| 276 | if choice.get("finish_reason"): |
| 277 | finish_reason = choice["finish_reason"] |
| 278 | |
| 279 | kv_params = chunk.get("kv_transfer_params") |
| 280 | if kv_params: |
| 281 | kv_params["remote_host"] = client.host |
| 282 | captured_kv = kv_params |
| 283 | if conversation_id: |
| 284 | kv_cache.put(conversation_id, kv_params) |
| 285 | |
| 286 | return ( |
| 287 | collected_text, |
| 288 | finish_reason, |
| 289 | captured_kv, |
| 290 | response_id or request_id, |
| 291 | model_name, |
| 292 | created, |
| 293 | ) |
| 294 | |
| 295 | |
| 296 | async def _stream_from_decode_sse( |
| 297 | client: ServiceClient, |
| 298 | endpoint: str, |
| 299 | req_data: dict[str, Any], |
| 300 | request_id: str, |
| 301 | conversation_id: str, |
| 302 | ): |
| 303 | """Yield SSE chunks from D to the client, capturing kv_transfer_params.""" |
| 304 | payload = req_data.copy() |
| 305 | payload["stream"] = True |
| 306 | |
| 307 | async with client.client.stream( |
| 308 | "POST", |
| 309 | endpoint, |
| 310 | json=payload, |
| 311 | headers=_make_headers(request_id), |
| 312 | ) as resp: |
| 313 | resp.raise_for_status() |
| 314 | async for line in resp.aiter_lines(): |
| 315 | if not line: |
| 316 | yield "\n" |
| 317 | continue |
| 318 | |
| 319 | if line.startswith("data: ") and line != "data: [DONE]": |
| 320 | try: |
| 321 | chunk = json.loads(line[6:]) |
| 322 | kv_params = chunk.get("kv_transfer_params") |
| 323 | if kv_params and conversation_id: |
| 324 | kv_params["remote_host"] = client.host |
| 325 | kv_cache.put(conversation_id, kv_params) |
| 326 | except json.JSONDecodeError: |
| 327 | pass |
| 328 | |
| 329 | yield line + "\n" |
| 330 | |
| 331 | |
| 332 | # FastAPI application |
| 333 | @asynccontextmanager |
| 334 | async def lifespan(app: FastAPI): |
| 335 | """Initialize HTTP clients for P and D instances.""" |
| 336 | app.state.prefill_clients: list[ServiceClient] = [] |
| 337 | app.state.decode_clients: list[ServiceClient] = [] |
| 338 | |
| 339 | for i, (host, port) in enumerate(global_args.prefiller_instances): |
| 340 | app.state.prefill_clients.append( |
| 341 | ServiceClient( |
| 342 | client=httpx.AsyncClient( |
| 343 | timeout=None, |
| 344 | base_url=f"http://{host}:{port}/v1", |
| 345 | ), |
| 346 | host=host, |
| 347 | port=port, |
| 348 | id=i, |
| 349 | ) |
| 350 | ) |
| 351 | |
| 352 | for i, (host, port) in enumerate(global_args.decoder_instances): |
| 353 | app.state.decode_clients.append( |
| 354 | ServiceClient( |
| 355 | client=httpx.AsyncClient( |
| 356 | timeout=None, |
| 357 | base_url=f"http://{host}:{port}/v1", |
| 358 | ), |
| 359 | host=host, |
| 360 | port=port, |
| 361 | id=i, |
| 362 | ) |
| 363 | ) |
| 364 | |
| 365 | app.state.prefill_iter = itertools.cycle(range(len(app.state.prefill_clients))) |
| 366 | app.state.decode_iter = itertools.cycle(range(len(app.state.decode_clients))) |
| 367 | |
| 368 | logger.info( |
| 369 | "Ready: %d prefill, %d decode instances", |
| 370 | len(app.state.prefill_clients), |
| 371 | len(app.state.decode_clients), |
| 372 | ) |
| 373 | yield |
| 374 | |
| 375 | for sc in app.state.prefill_clients + app.state.decode_clients: |
| 376 | await sc.client.aclose() |
| 377 | |
| 378 | |
| 379 | app = FastAPI(title="Disaggregated P/D Proxy (Multi-turn)", lifespan=lifespan) |
| 380 | |
| 381 | |
| 382 | def _next_client(app_state, role: str) -> ServiceClient: |
| 383 | if role == "prefill": |
| 384 | return app_state.prefill_clients[next(app_state.prefill_iter)] |
| 385 | return app_state.decode_clients[next(app_state.decode_iter)] |
| 386 | |
| 387 | |
| 388 | # Request handler |
| 389 | async def _handle_request(api_path: str, request: Request): |
| 390 | """Core request handler for both /v1/chat/completions and /v1/completions.""" |
| 391 | req_data = await request.json() |
| 392 | request_id = str(uuid.uuid4()) |
| 393 | conversation_id: str = req_data.pop("conversation_id", "") |
| 394 | client_wants_stream = req_data.get("stream", False) |
| 395 | |
| 396 | if not conversation_id: |
| 397 | logger.warning( |
| 398 | "[%s] No conversation_id provided — KV cache reuse disabled " |
| 399 | "for this request. Add a 'conversation_id' field to enable " |
| 400 | "cross-turn KV sharing. When using " |
| 401 | "benchmarks/multi_turn/benchmark_serving_multi_turn.py, pass " |
| 402 | "--send-conversation-id (off by default).", |
| 403 | request_id, |
| 404 | ) |
| 405 | |
| 406 | # Step 1: Look up cached D blocks from the previous turn |
| 407 | cached_kv = kv_cache.get(conversation_id) if conversation_id else None |
| 408 | |
| 409 | if cached_kv: |
| 410 | # Tell P to read D's blocks (bidirectional transfer) |
| 411 | cached_kv["do_remote_decode"] = True |
| 412 | cached_kv["do_remote_prefill"] = False |
| 413 | req_data["kv_transfer_params"] = cached_kv |
| 414 | logger.info( |
| 415 | "[%s] conv=%s: sending D's cached blocks to P (remote_request_id=%s)", |
| 416 | request_id, |
| 417 | conversation_id, |
| 418 | cached_kv.get("remote_request_id"), |
| 419 | ) |
| 420 | else: |
| 421 | # No cached blocks — P recomputes from scratch |
| 422 | req_data["kv_transfer_params"] = { |
| 423 | "do_remote_decode": True, |
| 424 | "do_remote_prefill": False, |
| 425 | "remote_engine_id": None, |
| 426 | "remote_block_ids": None, |
| 427 | "remote_host": None, |
| 428 | "remote_port": None, |
| 429 | } |
| 430 | logger.info("[%s] conv=%s: cache MISS", request_id, conversation_id) |
| 431 | |
| 432 | # Step 2: Send to Prefill node (non-streaming, max_tokens=1) |
| 433 | prefill_client = _next_client(request.app.state, "prefill") |
| 434 | t0 = time.time() |
| 435 | prefill_resp = await _send_to_prefill( |
| 436 | prefill_client, |
| 437 | api_path, |
| 438 | req_data, |
| 439 | request_id, |
| 440 | ) |
| 441 | logger.info( |
| 442 | "[%s] Prefill done in %.0fms", |
| 443 | request_id, |
| 444 | (time.time() - t0) * 1000, |
| 445 | ) |
| 446 | |
| 447 | # Attach P's kv_transfer_params for D to read P's blocks |
| 448 | p_kv_params = prefill_resp.get("kv_transfer_params", {}) |
| 449 | if p_kv_params: |
| 450 | p_kv_params["remote_host"] = prefill_client.host |
| 451 | req_data["kv_transfer_params"] = p_kv_params |
| 452 | |
| 453 | # Step 3: Stream from Decode node, capturing kv_transfer_params |
| 454 | decode_client = _next_client(request.app.state, "decode") |
| 455 | |
| 456 | if client_wants_stream: |
| 457 | return StreamingResponse( |
| 458 | _stream_from_decode_sse( |
| 459 | decode_client, |
| 460 | api_path, |
| 461 | req_data, |
| 462 | request_id, |
| 463 | conversation_id, |
| 464 | ), |
| 465 | media_type="text/event-stream", |
| 466 | ) |
| 467 | |
| 468 | text, finish_reason, _, resp_id, model, created = await _stream_from_decode( |
| 469 | decode_client, |
| 470 | api_path, |
| 471 | req_data, |
| 472 | request_id, |
| 473 | conversation_id, |
| 474 | ) |
| 475 | |
| 476 | # Build OpenAI-compatible response |
| 477 | is_chat = "messages" in req_data |
| 478 | if is_chat: |
| 479 | body = { |
| 480 | "id": resp_id, |
| 481 | "object": "chat.completion", |
| 482 | "created": created or int(time.time()), |
| 483 | "model": model or req_data.get("model", ""), |
| 484 | "choices": [ |
| 485 | { |
| 486 | "index": 0, |
| 487 | "message": {"role": "assistant", "content": text}, |
| 488 | "finish_reason": finish_reason, |
| 489 | } |
| 490 | ], |
| 491 | "usage": None, |
| 492 | } |
| 493 | else: |
| 494 | body = { |
| 495 | "id": resp_id, |
| 496 | "object": "text_completion", |
| 497 | "created": created or int(time.time()), |
| 498 | "model": model or req_data.get("model", ""), |
| 499 | "choices": [ |
| 500 | { |
| 501 | "index": 0, |
| 502 | "text": text, |
| 503 | "logprobs": None, |
| 504 | "finish_reason": finish_reason, |
| 505 | } |
| 506 | ], |
| 507 | "usage": None, |
| 508 | } |
| 509 | return JSONResponse(content=body) |
| 510 | |
| 511 | |
| 512 | # Routes |
| 513 | @app.post("/v1/chat/completions") |
| 514 | async def chat_completions(request: Request): |
| 515 | return await _handle_request("/chat/completions", request) |
| 516 | |
| 517 | |
| 518 | @app.post("/v1/completions") |
| 519 | async def completions(request: Request): |
| 520 | return await _handle_request("/completions", request) |
| 521 | |
| 522 | |
| 523 | @app.get("/health") |
| 524 | async def health(): |
| 525 | evicted = kv_cache.evict_stale() |
| 526 | return { |
| 527 | "status": "ok", |
| 528 | "cached_conversations": kv_cache.size, |
| 529 | "evicted_stale": evicted, |
| 530 | } |
| 531 | |
| 532 | |
| 533 | # CLI |
| 534 | def parse_args() -> argparse.Namespace: |
| 535 | p = argparse.ArgumentParser( |
| 536 | description="Disaggregated P/D proxy with bidirectional KV transfer", |
| 537 | ) |
| 538 | p.add_argument("--host", default="0.0.0.0") |
| 539 | p.add_argument("--port", type=int, default=8000) |
| 540 | p.add_argument( |
| 541 | "--prefiller-host", |
| 542 | "--prefiller-hosts", |
| 543 | dest="prefiller_hosts", |
| 544 | nargs="+", |
| 545 | default=["localhost"], |
| 546 | ) |
| 547 | p.add_argument( |
| 548 | "--prefiller-port", |
| 549 | "--prefiller-ports", |
| 550 | dest="prefiller_ports", |
| 551 | type=int, |
| 552 | nargs="+", |
| 553 | default=[8100], |
| 554 | ) |
| 555 | p.add_argument( |
| 556 | "--decoder-host", |
| 557 | "--decoder-hosts", |
| 558 | dest="decoder_hosts", |
| 559 | nargs="+", |
| 560 | default=["localhost"], |
| 561 | ) |
| 562 | p.add_argument( |
| 563 | "--decoder-port", |
| 564 | "--decoder-ports", |
| 565 | dest="decoder_ports", |
| 566 | type=int, |
| 567 | nargs="+", |
| 568 | default=[8200], |
| 569 | ) |
| 570 | args = p.parse_args() |
| 571 | |
| 572 | if len(args.prefiller_hosts) != len(args.prefiller_ports): |
| 573 | p.error("Number of prefiller hosts must match ports") |
| 574 | if len(args.decoder_hosts) != len(args.decoder_ports): |
| 575 | p.error("Number of decoder hosts must match ports") |
| 576 | |
| 577 | args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports)) |
| 578 | args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports)) |
| 579 | return args |
| 580 | |
| 581 | |
| 582 | if __name__ == "__main__": |
| 583 | global global_args |
| 584 | global_args = parse_args() |
| 585 | |
| 586 | import uvicorn |
| 587 | |
| 588 | uvicorn.run(app, host=global_args.host, port=global_args.port) |