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>

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)