diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index f7b2d63..6783795 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -30,6 +30,8 @@ CACHE_HIT_ALPHA = 1.0 HEAVY_THRESHOLD = 20000 # default; overridden by --heavy-threshold OVERLOAD_FACTOR = 2.0 # default; overridden by --overload-factor MAX_OFFLOAD_INFLIGHT = 4 # cap concurrent P-role offloads +PREFILL_THROUGHPUT = 7000 # tokens/s per GPU (from H20 measurements) +RDMA_OVERHEAD_S = 2.0 # seconds of RDMA transfer + decode start overhead class InstanceState: @@ -275,9 +277,9 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h "t_proxy_recv": _time.monotonic(), } - # H4 cache-aware offload gate: only offload when C_s has significant cache - # Cold turn-1 HEAVY: stay co-located (no RDMA overhead) - # Cached turn-2+ HEAVY: offload to flexible D (C_s fast prefill + D decode) + # Runtime cost-model offload gate: compare co-located vs offload latency + # Co-located = queue(C_s) + prefill(new_tokens) + # Offload = queue(P) + prefill(P_new_tokens) + RDMA_overhead offload_enabled = getattr(global_args, 'offload', False) and len(combined_instances) >= 2 use_offload = False offload_reason = "offload_disabled" @@ -285,28 +287,48 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h if estimated_new >= HEAVY_THRESHOLD and offload_enabled: cache_ratio = cache_hit / max(input_length, 1) current_offloads = sum(c.active_p_offloads for c in combined_instances) - d_candidate = min((c for c in combined_instances if c is not best_inst), + # P candidate: least-loaded instance (excluding C_s) + p_candidate = min((c for c in combined_instances if c is not best_inst), key=lambda c: c.ongoing_tokens) + # D candidate: least-loaded excluding both C_s and P + remaining = [c for c in combined_instances if c is not best_inst and c is not p_candidate] + d_candidate = min(remaining, key=lambda c: c.ongoing_tokens) if remaining else p_candidate + + # Cost model: compare co-located vs offload expected latency + # Co-located: queue on C_s + prefill new tokens on C_s + cs_queue = best_inst.pending_prefill_tokens / PREFILL_THROUGHPUT + colocated_cost = cs_queue + estimated_new / PREFILL_THROUGHPUT + + # Offload: prefill on P (may or may not have cache) + RDMA + decode start + p_queue = p_candidate.pending_prefill_tokens / PREFILL_THROUGHPUT + p_cache_hit = p_candidate.estimate_cache_hit(token_ids) if token_ids else 0 + p_new_tokens = max(0, input_length - p_cache_hit) + offload_cost = p_queue + p_new_tokens / PREFILL_THROUGHPUT + RDMA_OVERHEAD_S + breakdown["cache_ratio"] = cache_ratio + breakdown["colocated_cost"] = round(colocated_cost, 2) + breakdown["offload_cost"] = round(offload_cost, 2) if current_offloads >= MAX_OFFLOAD_INFLIGHT: offload_reason = "cap_reached_%d" % current_offloads - elif cache_ratio >= 0.3: + elif offload_cost < colocated_cost: use_offload = True - offload_reason = "cached_offload_%.0f%%" % (cache_ratio * 100) + offload_reason = "cost_model_%.1fvs%.1f" % (offload_cost, colocated_cost) else: - offload_reason = "cold_colocated_%.0f%%" % (cache_ratio * 100) + offload_reason = "colocated_cheaper_%.1fvs%.1f" % (colocated_cost, offload_cost) if use_offload: - p_inst = best_inst + p_inst = p_candidate d_inst = d_candidate d_idx = combined_instances.index(d_inst) # Accounting: reserve both P and D immediately so router sees the load + p_new = max(0, input_length - p_inst.estimate_cache_hit(token_ids)) if token_ids else input_length p_inst.ongoing_tokens += input_length - p_inst.pending_prefill_tokens += estimated_new + p_inst.pending_prefill_tokens += p_new p_inst.num_requests += 1 p_inst.active_p_offloads += 1 + breakdown["p_new_tokens"] = p_new d_inst.ongoing_tokens += input_length d_inst.num_requests += 1 @@ -370,15 +392,13 @@ PREFILL_TIMEOUT_S = 120 # max seconds to wait for P-instance prefill async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length, p_inst, d_inst, breakdown): - """HEAVY request: prefill on p_inst (C_s), KV via Mooncake, decode on d_inst (D). + """HEAVY request: prefill on p_inst, KV via Mooncake, decode on d_inst. On prefill timeout/failure, falls back to co-located decode on d_inst. """ request_id = headers.get("X-Request-Id", "") estimated_new = breakdown.get("estimated_new_tokens", 0) - # V2: p_inst is C_s with cache, so pending_prefill_tokens was incremented - # by estimated_new (only new tokens), not full input_length. - p_prefill_release = estimated_new + p_prefill_release = breakdown.get("p_new_tokens", estimated_new) # Step 1: Await prefill on p_inst (ongoing_tokens already reserved by caller) breakdown["t_prefill_sent"] = _time.monotonic()