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) <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
86
scripts/compare_adaptive.py
Normal file
86
scripts/compare_adaptive.py
Normal file
@@ -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))
|
||||
Reference in New Issue
Block a user