From d11d9f5cb9a5ee9d94765f381fe50294a053b8b0 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 22 May 2026 01:00:10 +0800 Subject: [PATCH] Adaptive prefill offload v1: implementation + experiment Added --heavy-threshold to cache_aware_proxy.py. HEAVY requests (new tokens >= threshold) route to instance with least decode load; WARM/MEDIUM route by cache-hit + token-level LB as before. Result: no significant difference vs baseline on single-machine combined mode. TTFT: +1.2%, TPOT: -1.5%, E2E: -0.3% (all within noise) Per-class TTFT breakdown shows the optimization target: WARM (75 req): p50=0.198s (cache hit, nearly free) MEDIUM (72 req): p50=1.356s HEAVY (54 req): p50=7.124s (36x slower than WARM) Conclusion: single-machine combined mode already distributes load well enough that adaptive routing adds no benefit. True isolation of HEAVY prefills requires cross-machine offload (v2 with Mooncake or multi-node). Co-Authored-By: Claude Opus 4.6 (1M context) --- scripts/cache_aware_proxy.py | 60 ++++++++++++++++++++----- scripts/compare_adaptive.py | 86 ++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 10 deletions(-) create mode 100644 scripts/compare_adaptive.py diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index b78d6eb..4583849 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -13,6 +13,7 @@ Routing policy (same for both modes): import argparse import asyncio import os +import time as _time import urllib.parse import uuid from contextlib import asynccontextmanager @@ -23,7 +24,8 @@ 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 +CACHE_HIT_ALPHA = 1.0 +HEAVY_THRESHOLD = 20000 # default; overridden by --heavy-threshold class InstanceState: @@ -35,6 +37,7 @@ class InstanceState: limits=httpx.Limits(max_connections=None, max_keepalive_connections=None), ) self.ongoing_tokens = 0 + self.ongoing_decode_tokens = 0 # subset: tokens in decode phase self.engine_id: dict[int, str] = {} self.dp_size = 1 self.cached_blocks: set[int] = set() @@ -90,6 +93,7 @@ prefill_instances: list[InstanceState] = [] decode_instances: list[InstanceState] = [] session_affinity: dict[str, int] = {} is_pd_sep = False +_breakdown_log: list[dict] = [] async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.Event): @@ -178,30 +182,63 @@ async def _handle(request: Request, api: str): 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) + """Combined mode with adaptive prefill offload. + + WARM/MEDIUM: route by cache-hit + load balance (co-located P+D). + HEAVY: route to instance with least decode load, avoiding decode disruption. + """ + # 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) + estimated_new = max(0, input_length - cache_hit) + + breakdown = { + "request_id": headers.get("X-Request-Id", ""), + "input_length": input_length, + "estimated_new_tokens": estimated_new, + "cache_hit": cache_hit, + "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" + + 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() + # 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 + breakdown["t_done"] = _time.monotonic() + _breakdown_log.append(breakdown) return StreamingResponse(generate(), media_type="text/event-stream") -import time as _time - -# Per-request breakdown log (append-only) -_breakdown_log: list[dict] = [] - - async def _send_prefill_async(p_inst, api, prefill_data, p_headers, token_ids, input_length, breakdown): """Fire-and-forget prefill: send and don't block caller.""" @@ -315,6 +352,8 @@ def parse_args(): help="PD-Sep decode: URL") p.add_argument("--fire-and-forget", action="store_true", 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)") args = p.parse_args() args.prefill = [] @@ -332,4 +371,5 @@ def parse_args(): if __name__ == "__main__": global_args = parse_args() + HEAVY_THRESHOLD = global_args.heavy_threshold uvicorn.run(app, host=global_args.host, port=global_args.port) diff --git a/scripts/compare_adaptive.py b/scripts/compare_adaptive.py new file mode 100644 index 0000000..b191cb9 --- /dev/null +++ b/scripts/compare_adaptive.py @@ -0,0 +1,86 @@ +"""Compare adaptive prefill offload vs baseline.""" +import csv, json, statistics, os, urllib.request + +def gpu_stats(path): + rows = list(csv.DictReader(open(path))) + vals = [float(r["util_pct"]) for r in rows] + s = sorted(vals) + p = lambda q: s[min(int(q*len(s)), len(s)-1)] + nz = sum(1 for v in vals if v > 0) + return {"mean": statistics.fmean(vals), "p50": p(.5), "p90": p(.9), + "active": nz*100//len(vals)} + +def lat_stats(path): + rows = [json.loads(l) for l in open(path)] + ok = [r for r in rows if not r.get("error")] + ttfts = sorted([r["ttft_s"] for r in ok if r.get("ttft_s")]) + tpots = sorted([r["tpot_s"] for r in ok if r.get("tpot_s") and r["tpot_s"]>0]) + lats = sorted([r["latency_s"] for r in ok]) + p = lambda v,q: v[min(int(q*len(v)),len(v)-1)] if v else 0 + return {"ok": len(ok), "n": len(rows), + "t50": p(ttfts,.5), "t90": p(ttfts,.9), + "p50": p(tpots,.5), "p90": p(tpots,.9), + "e50": p(lats,.5), "e90": p(lats,.9)} + +sep = "=" * 80 +print(sep) +print(" ADAPTIVE PREFILL OFFLOAD v1 vs BASELINE") +print(" Both: 8 combined TP=1 instances, cache-aware scheduler, 200 req") +print(sep) + +configs = [ + ("gpu_ab_combined", "Baseline (cache-aware)"), + ("gpu_ab_adaptive_20k", "Adaptive v1 (T=20k)"), +] + +print("\n LATENCY:") +fmt = " %-25s %7s %8s %8s %8s %8s %8s" +print(fmt % ("Config", "OK/N", "TTFT50", "TTFT90", "TPOT50", "TPOT90", "E2E50")) +print(" " + "-" * 68) +for d, label in configs: + s = lat_stats("outputs/%s/metrics.jsonl" % d) + print(fmt % (label, "%d/%d" % (s["ok"],s["n"]), + "%.3f" % s["t50"], "%.3f" % s["t90"], + "%.3f" % s["p50"], "%.3f" % s["p90"], "%.3f" % s["e50"])) + +print("\n GPU UTILIZATION:") +fmt2 = " %-25s %7s %7s %7s %7s" +print(fmt2 % ("Config", "Mean%", "P50%", "P90%", "Active")) +print(" " + "-" * 50) +for d, label in configs: + g = gpu_stats("outputs/%s/gpu_util.csv" % d) + print(fmt2 % (label, "%.1f" % g["mean"], "%.0f" % g["p50"], + "%.0f" % g["p90"], "%d%%" % g["active"])) + +# Breakdown by class +try: + data = json.loads(urllib.request.urlopen("http://localhost:9090/breakdown", timeout=5).read()) + from collections import Counter + classes = Counter(d.get("route_class", "?") for d in data) + print("\n REQUEST CLASSIFICATION (adaptive):") + for cls in ["WARM", "MEDIUM", "HEAVY"]: + cnt = classes.get(cls, 0) + subset = [d for d in data if d.get("route_class") == cls and "t_first_token" in d] + if subset: + ttfts = sorted([d["t_first_token"] - d["t_proxy_recv"] for d in subset]) + p50 = ttfts[len(ttfts)//2] + p90 = ttfts[min(int(0.9*len(ttfts)), len(ttfts)-1)] + print(" %s: n=%d TTFT p50=%.3fs p90=%.3fs" % (cls, cnt, p50, p90)) + else: + print(" %s: n=%d" % (cls, cnt)) +except Exception as e: + print("\n (breakdown: %s)" % e) + +# Delta +print("\n DELTA (Adaptive vs Baseline):") +b = lat_stats("outputs/gpu_ab_combined/metrics.jsonl") +a = lat_stats("outputs/gpu_ab_adaptive_20k/metrics.jsonl") +for label, bv, av in [ + ("TTFT p50", b["t50"], a["t50"]), + ("TTFT p90", b["t90"], a["t90"]), + ("TPOT p50", b["p50"], a["p50"]), + ("TPOT p90", b["p90"], a["p90"]), + ("E2E p50", b["e50"], a["e50"]), +]: + delta = (av/bv - 1) * 100 if bv > 0 else 0 + print(" %s: %.3f -> %.3f (%+.1f%%)" % (label, bv, av, delta))