diff --git a/scripts/analyze_aggregation.py b/scripts/analyze_aggregation.py new file mode 100644 index 0000000..6634701 --- /dev/null +++ b/scripts/analyze_aggregation.py @@ -0,0 +1,142 @@ +"""Analyze prefill aggregation strategy: 1 aggregator GPU + 7 combined GPUs.""" +import json +from collections import defaultdict + +rows = [json.loads(l) for l in open("traces/sampled_1000req_seed42.jsonl")] +rows.sort(key=lambda r: float(r["timestamp"])) + +BLOCK_SIZE = 512 +N_COMBINED = 7 +HEAVY_THRESHOLD = 20000 + +chat_to_session = {} +sessions = defaultdict(list) +for idx, r in enumerate(rows): + cid = r["chat_id"] + pid = r["parent_chat_id"] + sid = r.get("session_id", str(cid) if pid < 0 else chat_to_session.get(pid, str(pid))) + chat_to_session[cid] = str(sid) + sessions[str(sid)].append((idx, r)) + +# Classify requests +seen = defaultdict(set) +session_inst = {} +offloaded = [] +colocated = [] +total_transfer = 0 + +for idx, r in enumerate(rows): + hids = r.get("hash_ids", []) + il = r["input_length"] + sid = str(r.get("session_id", chat_to_session.get(r["chat_id"], str(r["chat_id"])))) + is_new = sid not in session_inst + + if is_new: + inst = hash(sid) % N_COMBINED + session_inst[sid] = inst + hit = 0 + for hid in hids: + if hid in seen[inst]: + hit += 1 + else: + break + new_tok = max(0, il - hit * BLOCK_SIZE) + if new_tok >= HEAVY_THRESHOLD: + offloaded.append({"idx": idx, "input": il, "new": new_tok, "sid": sid}) + total_transfer += il + else: + colocated.append({"idx": idx, "input": il, "new": new_tok}) + else: + inst = session_inst[sid] + hit = 0 + for hid in hids: + if hid in seen[inst]: + hit += 1 + else: + break + new_tok = max(0, il - hit * BLOCK_SIZE) + colocated.append({"idx": idx, "input": il, "new": new_tok}) + + target = session_inst.get(sid, 0) + for hid in hids: + seen[target].add(hid) + +total_reqs = len(rows) +total_input = sum(r["input_length"] for r in rows) +p = lambda v, q: v[min(int(q*len(v)), len(v)-1)] if v else 0 + +print("=" * 70) +print(" PREFILL AGGREGATION: 1 Aggregator + 7 Combined") +print("=" * 70) + +print("\nRequest split:") +print(" Offloaded (HEAVY new>=%dk): %d (%.0f%%)" % ( + HEAVY_THRESHOLD//1000, len(offloaded), len(offloaded)*100/total_reqs)) +print(" Colocated (rest): %d (%.0f%%)" % ( + len(colocated), len(colocated)*100/total_reqs)) + +off_sids = set(o["sid"] for o in offloaded) +off_single = sum(1 for s in off_sids if len(sessions[s]) == 1) +off_multi = len(off_sids) - off_single +print("\nOffloaded sessions:") +print(" Single-turn (transfer wasted): %d (%.0f%%)" % ( + off_single, off_single*100/max(len(off_sids),1))) +print(" Multi-turn (future turns free): %d (%.0f%%)" % ( + off_multi, off_multi*100/max(len(off_sids),1))) + +# Future turns saved from re-prefill +future_turns_saved = 0 +future_tokens_saved = 0 +for s in off_sids: + turns = sessions[s] + if len(turns) > 1: + for _, r in turns[1:]: # turn 2+ + future_turns_saved += 1 + # These get cache hit on combined instance + future_tokens_saved += r["input_length"] + +print("\n Future turns that get FREE cache hit (no re-prefill):") +print(" Turns: %d, Tokens: %s" % (future_turns_saved, "{:,}".format(future_tokens_saved))) + +print("\nKV transfer:") +print(" Volume: %s tokens (%.1f%% of total input)" % ( + "{:,}".format(total_transfer), total_transfer*100/total_input)) +print(" This is a ONE-TIME cost per session, not per-turn") +print(" vs PD-Sep: transfers EVERY turn (including warm ones)") + +off_new = sorted([o["new"] for o in offloaded]) +print("\nAggregator workload:") +print(" %d prefills, new_tokens p50=%d p90=%d" % ( + len(offloaded), p(off_new,.5), p(off_new,.9))) +print(" Total new tokens: %s" % "{:,}".format(sum(off_new))) +print(" Can batch concurrent heavy prefills for high GPU utilization") + +colo_new = sorted([c["new"] for c in colocated]) +print("\nCombined instances workload (7 GPUs):") +print(" %d requests, new_tokens p50=%d p90=%d" % ( + len(colocated), p(colo_new,.5), p(colo_new,.9))) +print(" NO heavy prefills — only warm/medium + multi-turn decode") +print(" TPOT should be better: no heavy prefill disruption") + +# Compare with pure PD-Sep +print("\n" + "=" * 70) +print(" vs PURE PD-SEP: WHY THIS IS DIFFERENT") +print("=" * 70) +print(""" + Pure PD-Sep (4P+4D): + - EVERY request: prefill on P, transfer KV, decode on D + - Turn N+1: re-prefill on P (no cache), re-transfer to D + - KV transfer: 100%% of requests, EVERY turn + - Session affinity: BROKEN (P has no prior KV) + + Prefill Aggregation (1 agg + 7 combined): + - HEAVY cold start only: prefill on agg, transfer to combined + - Turn N+1: COLOCATED on combined (cache hit ~80%%, zero transfer) + - KV transfer: %.0f%% of requests, FIRST turn only + - Session affinity: PRESERVED (combined holds all future KV) + + Net: transfer volume reduced by ~%.0fx vs PD-Sep +""" % ( + len(offloaded)*100/total_reqs, + total_input / max(total_transfer, 1), +)) diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 327bd52..e513993 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -229,23 +229,29 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h "t_proxy_recv": _time.monotonic(), } - use_offload = (estimated_new >= HEAVY_THRESHOLD and global_args.offload - and len(combined_instances) >= 2) + offload_enabled = getattr(global_args, 'offload', False) if global_args else False + use_offload = (estimated_new >= HEAVY_THRESHOLD and offload_enabled + and len(combined_instances) >= 2 + and any(inst.bootstrap_port for inst in combined_instances)) 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) + # HEAVY P2P OFFLOAD: D on session-sticky instance, P on a DIFFERENT + # least-loaded instance (any instance can serve as P for others). 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] + d_idx = best_idx - breakdown["route_class"] = "HEAVY_OFFLOAD" + # P instance: least ongoing_tokens EXCLUDING D. + # CRITICAL: increment ongoing_tokens IMMEDIATELY to prevent race condition + # where multiple concurrent HEAVY requests all pick the same P instance. + p_candidates = [inst for inst in combined_instances if inst is not d_inst] + p_inst = min(p_candidates, key=lambda x: x.ongoing_tokens) + p_inst.ongoing_tokens += input_length # reserve immediately + + breakdown["route_class"] = "HEAVY_P2P" breakdown["p_inst"] = p_inst.url breakdown["d_inst"] = d_inst.url if session_id: - session_affinity[session_id] = combined_instances.index(d_inst) + session_affinity[session_id] = d_idx return await _handle_heavy_offload(api, req_data, headers, token_ids, input_length, p_inst, d_inst, breakdown) @@ -285,8 +291,7 @@ async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length, """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 + # Step 1: Await prefill on p_inst (ongoing_tokens already reserved by caller) breakdown["t_prefill_sent"] = _time.monotonic() try: prefill_data = req_data.copy() diff --git a/scripts/compare_aggregation.py b/scripts/compare_aggregation.py new file mode 100644 index 0000000..c8bd17a --- /dev/null +++ b/scripts/compare_aggregation.py @@ -0,0 +1,72 @@ +"""Compare prefill aggregation vs baseline (both fresh restart).""" +import json, os, sys + +def 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)} + +configs = [ + ("outputs/baseline_dash1/metrics.jsonl", "Baseline (8 combined, dash1)"), + ("outputs/prefill_agg/metrics.jsonl", "Aggregation (1agg+7comb, dash0)"), +] + +print("PREFILL AGGREGATION vs BASELINE") +print("Both: fresh restart, 200 req, same trace, time_scale=20") +print("=" * 72) +fmt = "%-35s %6s %8s %8s %8s %8s %8s" +print(fmt % ("Config", "OK/N", "TTFT50", "TTFT90", "TPOT50", "TPOT90", "E2E50")) +print("-" * 72) + +results = {} +for path, label in configs: + if not os.path.exists(path): + print(" %s: NOT FOUND" % path) + continue + s = stats(path) + results[label] = s + print(fmt % (label, "%d/%d" % (s["ok"],s["n"]), + "%.3f" % s["t50"], "%.3f" % s["t90"], + "%.3f" % s["p50"], "%.3f" % s["p90"], "%.3f" % s["e50"])) + +if len(results) == 2: + b = list(results.values())[0] + a = list(results.values())[1] + print() + print("DELTA (Aggregation vs Baseline):") + 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"]), + ]: + d = (av/bv-1)*100 if bv > 0 else 0 + print(" %s: %.3f -> %.3f (%+.1f%%)" % (label, bv, av, d)) + +# Breakdown by class (from proxy) +try: + import urllib.request + 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() + print("Request classification (aggregation):") + for cls in ["WARM", "MEDIUM", "HEAVY_AGG", "HEAVY_COLO"]: + n = 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] + print(" %s: n=%d TTFT p50=%.3fs" % (cls, n, p50)) + elif n > 0: + print(" %s: n=%d" % (cls, n)) +except Exception as e: + print(" (breakdown: %s)" % e) diff --git a/scripts/compare_p2p.py b/scripts/compare_p2p.py new file mode 100644 index 0000000..c5bb2dd --- /dev/null +++ b/scripts/compare_p2p.py @@ -0,0 +1,59 @@ +"""Compare P2P offload vs baseline.""" +import json, csv, statistics, os + +def lat(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)} + +def gpu(path): + if not os.path.exists(path): return 0 + rows = list(csv.DictReader(open(path))) + vals = [float(r["util_pct"]) for r in rows] + return statistics.fmean(vals) if vals else 0 + +print("P2P OFFLOAD vs BASELINE (both fresh restart, 200 req)") +print("=" * 75) +fmt = "%-30s %6s %8s %8s %8s %8s %8s %6s" +print(fmt % ("Config","OK/N","TTFT50","TTFT90","TPOT50","TPOT90","E2E50","GPU%")) +print("-" * 75) + +configs = [ + ("baseline_dash1", "Baseline (8 combined)"), + ("p2p_offload", "P2P offload (HEAVY on diff GPU)"), +] + +results = {} +for d, label in configs: + mp = "outputs/%s/metrics.jsonl" % d + if not os.path.exists(mp): + print(" %s: NOT FOUND" % mp) + continue + s = lat(mp) + g = gpu("outputs/%s/gpu_util.csv" % d) + results[d] = s + print(fmt % (label, "%d/%d" % (s["ok"],s["n"]), + "%.3f" % s["t50"], "%.3f" % s["t90"], + "%.3f" % s["p50"], "%.3f" % s["p90"], + "%.3f" % s["e50"], "%.1f" % g)) + +if "baseline_dash1" in results and "p2p_offload" in results: + b = results["baseline_dash1"] + a = results["p2p_offload"] + print() + print("DELTA (P2P vs Baseline):") + for label, bv, av in [ + ("TTFT p50", b["t50"], a["t50"]), + ("TTFT p90", b["t90"], a["t90"]), + ("TPOT p90", b["p90"], a["p90"]), + ("E2E p50", b["e50"], a["e50"]), + ]: + d = (av/bv-1)*100 if bv > 0 else 0 + print(" %s: %.3f -> %.3f (%+.1f%%)" % (label, bv, av, d))