Replace static offload gate with runtime cost model

Old gate: cache_ratio >= 0.3 (static, only 14% of HEAVY triggered)
New gate: offload when offload_cost < colocated_cost, where:
  colocated_cost = queue(C_s) + prefill(new_tokens)
  offload_cost = queue(P_idle) + prefill(P_tokens) + RDMA_overhead

Key changes:
- P is now least-loaded instance (not session-sticky C_s)
- Gate considers C_s queue depth dynamically
- Crossover: offload wins when C_s queue >= 38k tokens (~5.4s)
- Cold HEAVY requests CAN be offloaded if C_s is busy enough
- P accounting uses P's actual cache hit, not C_s's

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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
+    # Runtime cost-model offload gate: compare co-located vs offload latency
-    # Cold turn-1 HEAVY: stay co-located (no RDMA overhead)
+    # Co-located = queue(C_s) + prefill(new_tokens)
-    # Cached turn-2+ HEAVY: offload to flexible D (C_s fast prefill + D decode)
+    # 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
+    p_prefill_release = breakdown.get("p_new_tokens", estimated_new)
-    # by estimated_new (only new tokens), not full input_length.
-    p_prefill_release = estimated_new
 
     # Step 1: Await prefill on p_inst (ongoing_tokens already reserved by caller)
     breakdown["t_prefill_sent"] = _time.monotonic()