From 2fee35562687b9f71ac1c5c57665fe8eaefa9b1f Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 22 May 2026 10:14:10 +0800 Subject: [PATCH] Adaptive v2 (selective Mooncake offload): worse than baseline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implemented --offload mode: HEAVY requests (>20k new tokens) get P on least-loaded instance, KV via Mooncake RDMA, D on session-sticky instance. WARM/MEDIUM stay co-located (no KV transfer). All 8 instances run kv_both. Result (200 req, same instances, fresh restart): Baseline (no offload): TTFT=1.073 TPOT90=0.074 E2E=5.086 Offload HEAVY: TTFT=1.462 TPOT90=0.077 E2E=6.847 Delta: +36% +4% +35% Conclusion: even selective KV transfer (only 44% of requests) adds more overhead than the isolation benefit provides. On single-machine 8 GPU, PD-combined with hybrid routing is strictly optimal. No form of KV transfer — full PD-sep, selective offload, or otherwise — improves over co-located serving for this workload. Co-Authored-By: Claude Opus 4.6 (1M context) --- scripts/cache_aware_proxy.py | 151 ++++++++++++++++++++++++++++------- 1 file changed, 122 insertions(+), 29 deletions(-) diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 012c179..327bd52 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -148,10 +148,15 @@ async def lifespan(app: FastAPI): 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") + bp_list = [int(p) for p in global_args.bootstrap_ports.split(",") if p.strip()] if global_args.bootstrap_ports else [] + for i, url in enumerate(global_args.combined): + bp = bp_list[i] if i < len(bp_list) else None + combined_instances.append(InstanceState(url, bp)) + if global_args.offload and bp_list: + await init_prefill_bootstrap(combined_instances, app.state.ready) + else: + app.state.ready.set() + print(f"Combined mode: {len(combined_instances)} instances, offload={'ON' if global_args.offload else 'OFF'}") else: is_pd_sep = True for url, bp in global_args.prefill: @@ -204,12 +209,13 @@ async def _handle(request: Request, api: str): async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers): - """Combined mode with adaptive prefill offload. + """Combined mode with adaptive prefill offload (v2). - WARM/MEDIUM: route by cache-hit + load balance (co-located P+D). - HEAVY: route to instance with least decode load, avoiding decode disruption. + WARM/MEDIUM: route to best instance, co-located P+D (no KV transfer). + HEAVY (kv_both mode): P on least-loaded instance, KV via Mooncake, D on + session-sticky instance. Only works if instances have kv_role=kv_both. + Falls back to co-located if --no-offload or instances lack Mooncake. """ - # Estimate new tokens after cache best_inst, best_idx = pick_instance(combined_instances, token_ids, session_id, input_length, session_affinity) cache_hit = best_inst.estimate_cache_hit(token_ids) @@ -223,42 +229,125 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h "t_proxy_recv": _time.monotonic(), } - if estimated_new >= HEAVY_THRESHOLD: - # HEAVY: pick instance with least ongoing_decode_tokens - # This avoids sending heavy prefill to an instance busy decoding - inst = min(combined_instances, key=lambda x: x.ongoing_decode_tokens) - idx = combined_instances.index(inst) - breakdown["route_class"] = "HEAVY" - if session_id: - session_affinity[session_id] = idx - else: - inst = best_inst - idx = best_idx - breakdown["route_class"] = "WARM" if estimated_new < 5000 else "MEDIUM" + use_offload = (estimated_new >= HEAVY_THRESHOLD and global_args.offload + and len(combined_instances) >= 2) - breakdown["routed_to"] = inst.url - inst.ongoing_tokens += input_length + if use_offload: + # HEAVY with offload: P on least-loaded, D on session-sticky (best_inst) + p_inst = min(combined_instances, key=lambda x: x.ongoing_tokens) + d_inst = best_inst + if p_inst is d_inst: + # Pick second-least-loaded for P + sorted_by_load = sorted(combined_instances, key=lambda x: x.ongoing_tokens) + p_inst = sorted_by_load[0] if sorted_by_load[0] is not d_inst else sorted_by_load[1] + + breakdown["route_class"] = "HEAVY_OFFLOAD" + breakdown["p_inst"] = p_inst.url + breakdown["d_inst"] = d_inst.url + if session_id: + session_affinity[session_id] = combined_instances.index(d_inst) + + return await _handle_heavy_offload(api, req_data, headers, token_ids, + input_length, p_inst, d_inst, breakdown) + else: + if estimated_new >= HEAVY_THRESHOLD: + breakdown["route_class"] = "HEAVY_COLO" + else: + breakdown["route_class"] = "WARM" if estimated_new < 5000 else "MEDIUM" + + inst = best_inst + breakdown["routed_to"] = inst.url + inst.ongoing_tokens += input_length + + async def generate(): + first_token = True + try: + async with inst.client.stream("POST", api, json=req_data, headers=headers) as resp: + resp.raise_for_status() + inst.ongoing_decode_tokens += input_length + async for chunk in resp.aiter_bytes(): + if first_token: + breakdown["t_first_token"] = _time.monotonic() + first_token = False + yield chunk + inst.record_prefix(token_ids) + finally: + inst.ongoing_tokens -= input_length + inst.ongoing_decode_tokens -= input_length + breakdown["t_done"] = _time.monotonic() + _breakdown_log.append(breakdown) + + return StreamingResponse(generate(), media_type="text/event-stream") + + +async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length, + p_inst, d_inst, breakdown): + """HEAVY request: prefill on p_inst, KV via Mooncake, decode on d_inst.""" + request_id = headers.get("X-Request-Id", "") + + # Step 1: Await prefill on p_inst + p_inst.ongoing_tokens += input_length + breakdown["t_prefill_sent"] = _time.monotonic() + try: + prefill_data = req_data.copy() + prefill_data["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "transfer_id": "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) + breakdown["t_prefill_done"] = _time.monotonic() + except Exception as e: + breakdown["t_prefill_done"] = _time.monotonic() + breakdown["error"] = str(e) + _breakdown_log.append(breakdown) + raise HTTPException(status_code=502, detail="Prefill failed: %s" % e) + finally: + p_inst.ongoing_tokens -= input_length + + # Step 2: Stream decode on d_inst (pulls KV from Mooncake) + d_inst.ongoing_tokens += input_length + d_inst.ongoing_decode_tokens += input_length + breakdown["t_decode_sent"] = _time.monotonic() + + parsed = urllib.parse.urlparse(str(p_inst.client.base_url)) + bootstrap_addr = "http://%s:%s" % (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": "xfer-" + request_id, + } async def generate(): first_token = True try: - async with inst.client.stream("POST", api, json=req_data, headers=headers) as resp: + async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp: resp.raise_for_status() - # Once streaming starts, this instance is in "decode phase" - inst.ongoing_decode_tokens += input_length async for chunk in resp.aiter_bytes(): if first_token: breakdown["t_first_token"] = _time.monotonic() first_token = False yield chunk - inst.record_prefix(token_ids) finally: - inst.ongoing_tokens -= input_length - inst.ongoing_decode_tokens -= input_length + d_inst.ongoing_tokens -= input_length + d_inst.ongoing_decode_tokens -= input_length breakdown["t_done"] = _time.monotonic() _breakdown_log.append(breakdown) - return StreamingResponse(generate(), media_type="text/event-stream") + return StreamingResponse(generate(), media_type="application/json") async def _send_prefill_async(p_inst, api, prefill_data, p_headers, token_ids, @@ -376,6 +465,10 @@ def parse_args(): help="Send prefill async, don't await before decode") p.add_argument("--heavy-threshold", type=int, default=20000, help="New tokens threshold for HEAVY classification (adaptive offload)") + p.add_argument("--offload", action="store_true", + help="Enable Mooncake KV offload for HEAVY requests (requires kv_both instances)") + p.add_argument("--bootstrap-ports", type=str, default="", + help="Comma-separated bootstrap ports for combined instances (for offload mode)") args = p.parse_args() args.prefill = []