"""Unified cache-aware + token-level load-balanced global scheduler. Supports two modes: --combined URL [URL ...]: PD co-located instances (normal vLLM, no KV transfer) --prefill URL BP --decode URL: PD disaggregated instances (Mooncake KV transfer) Routing policy (same for both modes): score = ongoing_tokens / avg_ongoing - ALPHA * cache_hit_ratio Normalized load prevents "rich get richer"; cache bonus gives affinity. Session affinity: multi-turn sessions stick to same instance. """ import argparse import asyncio import os import urllib.parse import uuid from contextlib import asynccontextmanager import httpx import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse BLOCK_SIZE = 512 CACHE_HIT_ALPHA = 1.0 # weight for cache bonus in scoring class InstanceState: def __init__(self, url: str, bootstrap_port: int | None = None): self.url = url self.bootstrap_port = bootstrap_port self.client = httpx.AsyncClient( timeout=None, base_url=url, limits=httpx.Limits(max_connections=None, max_keepalive_connections=None), ) self.ongoing_tokens = 0 self.engine_id: dict[int, str] = {} self.dp_size = 1 self.cached_blocks: set[int] = set() def estimate_cache_hit(self, token_ids: list[int] | None) -> int: if not token_ids or len(token_ids) < BLOCK_SIZE: return 0 hit = 0 for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE): bh = hash(tuple(token_ids[i:i + BLOCK_SIZE])) if bh in self.cached_blocks: hit += BLOCK_SIZE else: break return hit def record_prefix(self, token_ids: list[int] | None): if not token_ids: return for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE): self.cached_blocks.add(hash(tuple(token_ids[i:i + BLOCK_SIZE]))) if len(self.cached_blocks) > 200000: self.cached_blocks = set(list(self.cached_blocks)[-100000:]) def pick_instance(instances: list[InstanceState], token_ids: list[int] | None, session_id: str | None, input_length: int, affinity: dict[str, int]) -> tuple[InstanceState, int]: """Normalized load - cache bonus scoring.""" if session_id and session_id in affinity: idx = affinity[session_id] if idx < len(instances): return instances[idx], idx avg_load = max(sum(i.ongoing_tokens for i in instances) / len(instances), 1.0) best_idx, best_score = 0, float("inf") for i, inst in enumerate(instances): cache_hit = inst.estimate_cache_hit(token_ids) cache_ratio = cache_hit / input_length if input_length > 0 else 0.0 score = inst.ongoing_tokens / avg_load - CACHE_HIT_ALPHA * cache_ratio if score < best_score: best_score = score best_idx = i if session_id: affinity[session_id] = best_idx return instances[best_idx], best_idx global_args = None combined_instances: list[InstanceState] = [] prefill_instances: list[InstanceState] = [] decode_instances: list[InstanceState] = [] session_affinity: dict[str, int] = {} is_pd_sep = False async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.Event): for inst in instances: if inst.bootstrap_port is None: continue while True: try: await inst.client.get("/health") except Exception: await asyncio.sleep(1) continue parsed = urllib.parse.urlparse(str(inst.client.base_url)) url = f"http://{parsed.hostname}:{inst.bootstrap_port}/query" resp = await inst.client.get(url) resp.raise_for_status() data = resp.json() for dp_rank, dp_entry in data.items(): inst.engine_id[int(dp_rank)] = dp_entry["engine_id"] inst.dp_size = len(data) print(f"Inited {inst.url} engine_ids={inst.engine_id}") break ready.set() @asynccontextmanager async def lifespan(app: FastAPI): global is_pd_sep app.state.ready = asyncio.Event() if global_args.combined: is_pd_sep = False for url in global_args.combined: combined_instances.append(InstanceState(url)) app.state.ready.set() print(f"Combined mode: {len(combined_instances)} instances") else: is_pd_sep = True for url, bp in global_args.prefill: prefill_instances.append(InstanceState(url, bp)) for url in global_args.decode: decode_instances.append(InstanceState(url)) await init_prefill_bootstrap(prefill_instances, app.state.ready) print(f"PD-Sep mode: {len(prefill_instances)}P + {len(decode_instances)}D") yield for inst in combined_instances + prefill_instances + decode_instances: await inst.client.aclose() app = FastAPI(lifespan=lifespan) @app.post("/v1/completions") async def handle_completions(request: Request): return await _handle(request, "/v1/completions") @app.post("/v1/chat/completions") async def handle_chat(request: Request): return await _handle(request, "/v1/chat/completions") async def _handle(request: Request, api: str): if not app.state.ready.is_set(): raise HTTPException(status_code=503, detail="Service Unavailable") req_data = await request.json() request_id = str(uuid.uuid4()) prompt = req_data.get("prompt") token_ids = prompt if isinstance(prompt, list) else None input_length = len(token_ids) if token_ids else 0 session_id = request.headers.get("X-Session-Id") headers = {"X-Request-Id": request_id} api_key = os.environ.get("OPENAI_API_KEY") if api_key: headers["Authorization"] = f"Bearer {api_key}" if is_pd_sep: return await _handle_pd_sep(api, req_data, request_id, token_ids, input_length, session_id, headers) else: return await _handle_combined(api, req_data, token_ids, input_length, session_id, headers) async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers): """Combined mode: route to best instance, send normal request.""" inst, idx = pick_instance(combined_instances, token_ids, session_id, input_length, session_affinity) inst.ongoing_tokens += input_length async def generate(): try: async with inst.client.stream("POST", api, json=req_data, headers=headers) as resp: resp.raise_for_status() async for chunk in resp.aiter_bytes(): yield chunk inst.record_prefix(token_ids) finally: inst.ongoing_tokens -= input_length return StreamingResponse(generate(), media_type="text/event-stream") async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length, session_id, headers): """PD-Sep mode: await prefill, then stream decode.""" p_inst, _ = pick_instance(prefill_instances, token_ids, session_id, input_length, session_affinity) d_inst = min(decode_instances, key=lambda x: x.ongoing_tokens) # Await prefill p_inst.ongoing_tokens += input_length try: prefill_data = req_data.copy() prefill_data["kv_transfer_params"] = { "do_remote_decode": True, "do_remote_prefill": False, "transfer_id": f"xfer-{request_id}", } prefill_data["stream"] = False prefill_data["max_tokens"] = 1 prefill_data.pop("max_completion_tokens", None) prefill_data.pop("stream_options", None) p_headers = {**headers, "X-data-parallel-rank": "0"} resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers) resp.raise_for_status() await resp.aclose() p_inst.record_prefix(token_ids) except Exception as e: raise HTTPException(status_code=502, detail=f"Prefill failed: {e}") finally: p_inst.ongoing_tokens -= input_length # Stream decode d_inst.ongoing_tokens += input_length parsed = urllib.parse.urlparse(str(p_inst.client.base_url)) bootstrap_addr = f"http://{parsed.hostname}:{p_inst.bootstrap_port}" decode_data = req_data.copy() decode_data["kv_transfer_params"] = { "do_remote_decode": False, "do_remote_prefill": True, "remote_bootstrap_addr": bootstrap_addr, "remote_engine_id": p_inst.engine_id.get(0, ""), "transfer_id": f"xfer-{request_id}", } async def generate(): try: async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp: resp.raise_for_status() async for chunk in resp.aiter_bytes(): yield chunk finally: d_inst.ongoing_tokens -= input_length return StreamingResponse(generate(), media_type="application/json") def parse_args(): p = argparse.ArgumentParser(description="Unified cache-aware global scheduler") p.add_argument("--port", type=int, default=8000) p.add_argument("--host", type=str, default="0.0.0.0") p.add_argument("--combined", nargs="+", help="Combined mode: list of instance URLs") p.add_argument("--prefill", nargs="+", action="append", dest="prefill_raw", help="PD-Sep prefill: URL [bootstrap_port]") p.add_argument("--decode", nargs=1, action="append", dest="decode_raw", help="PD-Sep decode: URL") args = p.parse_args() args.prefill = [] if args.prefill_raw: for entry in args.prefill_raw: url = entry[0] bp = int(entry[1]) if len(entry) > 1 and entry[1].lower() != "none" else None args.prefill.append((url, bp)) args.decode = [e[0] for e in (args.decode_raw or [])] if not args.combined and not args.prefill: p.error("Must specify either --combined or --prefill/--decode") return args if __name__ == "__main__": global_args = parse_args() uvicorn.run(app, host=global_args.host, port=global_args.port)