Fix multi-turn replay fidelity: track realized output tokens across all components

The replayer and proxy were building multi-turn prompts from trace tokens,
but the model generates different output tokens. Subsequent turns had wrong
prefix tokens, causing cache misses and invalid experimental measurements.

- replay.py: min_tokens=max_tokens for deterministic length, return_token_ids
  to capture actual output, _apply_realized_prefix for next-turn correction
- proxy: extract output token_ids from SSE, record prompt+output as realized
  prefix in shadow cache, extract _handle_local_request to deduplicate
- bench.sh/launch_elastic_p2p.sh: default elastic mode to unified policy
- mooncake_connector: only send prompt blocks (not stale output blocks),
  track failed_recving_block_ids for error recovery

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

replayer/replay.py

@@ -84,6 +84,12 @@ class _SessionState:
     metrics: list[RequestMetrics] = field(default_factory=list)
 
 
+@dataclass
+class _DispatchResult:
+    metric: RequestMetrics
+    output_token_ids: list[int]
+
+
 _endpoint_counter = 0
 
 
@@ -103,14 +109,17 @@ async def _dispatch_request(
     req: TraceRequest,
     prompt_token_ids: list[int],
     sem: asyncio.Semaphore,
-) -> RequestMetrics:
+) -> _DispatchResult:
     """Send one request via /v1/completions (streaming) and collect metrics."""
     endpoint = _pick_endpoint(config)
+    target_output_tokens = max(1, req.output_length)
     payload = {
         "model": config.model_name,
         "prompt": prompt_token_ids,
-        "max_tokens": max(1, req.output_length),
+        "max_tokens": target_output_tokens,
+        "min_tokens": target_output_tokens,
         "temperature": 0,
+        "return_token_ids": True,
         "stream": True,
         "stream_options": {"include_usage": True},
     }
@@ -122,6 +131,7 @@ async def _dispatch_request(
     finish_reason = None
     err = None
     token_times: list[float] = []
+    output_token_ids: list[int] = []
 
     req_headers = {"X-Session-Id": req.session_id}
 
@@ -146,14 +156,24 @@ async def _dispatch_request(
                     except json.JSONDecodeError:
                         continue
 
-                    now = time.perf_counter()
-                    if ttft_s is None:
-                        ttft_s = now - start
-
                     choices = chunk.get("choices", [])
                     if choices:
+                        now = time.perf_counter()
                         delta = choices[0].get("text", "")
-                        if delta:
+                        chunk_token_ids = choices[0].get("token_ids")
+                        if isinstance(chunk_token_ids, list):
+                            clean_ids = [
+                                int(t) for t in chunk_token_ids
+                                if isinstance(t, int)
+                            ]
+                            if clean_ids:
+                                if ttft_s is None:
+                                    ttft_s = now - start
+                                output_token_ids.extend(clean_ids)
+                                token_times.extend([now] * len(clean_ids))
+                        elif delta:
+                            if ttft_s is None:
+                                ttft_s = now - start
                             token_times.append(now)
                         fr = choices[0].get("finish_reason")
                         if fr:
@@ -168,8 +188,15 @@ async def _dispatch_request(
 
     end = time.perf_counter()
     e2e = end - start
-    if n_output == 0 and token_times:
+    if output_token_ids:
+        n_output = len(output_token_ids)
+    elif n_output == 0 and token_times:
         n_output = len(token_times)
+    if err is None and n_output != target_output_tokens:
+        err = (
+            "output_token_mismatch "
+            f"requested={target_output_tokens} actual={n_output}"
+        )
 
     tpot = 0.0
     if len(token_times) > 1:
@@ -177,26 +204,42 @@ async def _dispatch_request(
                        for i in range(len(token_times) - 1)]
         tpot = sum(inter_token) / len(inter_token)
 
-    return RequestMetrics(
-        request_id=req.request_id,
-        session_id=req.session_id,
-        turn_id=req.turn_id,
-        trace_timestamp_s=req.timestamp_s,
-        input_length=req.input_length,
-        output_length=req.output_length,
-        request_type=req.request_type,
-        effective_input_length=len(prompt_token_ids),
-        cached_tokens=cached_tokens,
-        latency_s=e2e,
-        ttft_s=ttft_s,
-        tpot_s=tpot,
-        actual_output_tokens=n_output,
-        requested_output_tokens=req.output_length,
-        finish_reason=finish_reason,
-        error=err,
+    return _DispatchResult(
+        metric=RequestMetrics(
+            request_id=req.request_id,
+            session_id=req.session_id,
+            turn_id=req.turn_id,
+            trace_timestamp_s=req.timestamp_s,
+            input_length=req.input_length,
+            output_length=req.output_length,
+            request_type=req.request_type,
+            effective_input_length=len(prompt_token_ids),
+            cached_tokens=cached_tokens,
+            latency_s=e2e,
+            ttft_s=ttft_s,
+            tpot_s=tpot,
+            actual_output_tokens=n_output,
+            requested_output_tokens=req.output_length,
+            finish_reason=finish_reason,
+            error=err,
+        ),
+        output_token_ids=output_token_ids,
     )
 
 
+def _apply_realized_prefix(
+    prompt_token_ids: list[int],
+    realized_context: list[int],
+) -> list[int]:
+    """Replace the reusable session prefix with engine-realized tokens."""
+    if not realized_context:
+        return prompt_token_ids
+    out = prompt_token_ids.copy()
+    n = min(len(out), len(realized_context))
+    out[:n] = realized_context[:n]
+    return out
+
+
 def _extract_cached_tokens(usage: dict) -> int:
     ct = 0
     details = usage.get("prompt_tokens_details")
@@ -220,6 +263,7 @@ async def _run_session(
 ) -> list[RequestMetrics]:
     if session_sem is not None:
         await session_sem.acquire()
+    realized_context: list[int] = []
     try:
         for req in state.turns:
             # Wait until this request's trace timestamp
@@ -228,13 +272,19 @@ async def _run_session(
             if elapsed < target_wall:
                 await asyncio.sleep(target_wall - elapsed)
 
-            token_ids = _build_prompt_token_ids(req)
-            metric = await _dispatch_request(
+            token_ids = _apply_realized_prefix(
+                _build_prompt_token_ids(req),
+                realized_context,
+            )
+            result = await _dispatch_request(
                 client=client, config=config, req=req,
                 prompt_token_ids=token_ids, sem=request_sem,
             )
+            metric = result.metric
             state.metrics.append(metric)
             await sink.append(metric)
+            if metric.error is None:
+                realized_context = token_ids + result.output_token_ids
     finally:
         if session_sem is not None:
             session_sem.release()

scripts/bench.sh

@@ -26,7 +26,8 @@ TRACE="${TRACE:-$PROJECT_DIR/traces/w600_r0.0015_st30.jsonl}"
 # Defaults
 TAG=""
 MODE="baseline"     # baseline | elastic
-POLICY="linear"     # linear | lmetric
+POLICY="linear"     # linear | lmetric | unified
+POLICY_SET=false
 N_INSTANCES=8
 BASE_PORT=8000
 PROXY_PORT=9090
@@ -44,7 +45,7 @@ while [[ $# -gt 0 ]]; do
     case "$1" in
         --tag) TAG="$2"; shift 2 ;;
         --mode) MODE="$2"; shift 2 ;;
-        --policy) POLICY="$2"; shift 2 ;;
+        --policy) POLICY="$2"; POLICY_SET=true; shift 2 ;;
         --instances) N_INSTANCES="$2"; shift 2 ;;
         --requests) REQUESTS="$2"; shift 2 ;;
         --trace) TRACE="$2"; shift 2 ;;
@@ -60,11 +61,15 @@ while [[ $# -gt 0 ]]; do
 done
 
 if [ -z "$TAG" ]; then
-    echo "Usage: bench.sh --tag NAME --mode {baseline|elastic} [--instances N] [--policy {linear|lmetric}] [--requests N]"
+    echo "Usage: bench.sh --tag NAME --mode {baseline|elastic} [--instances N] [--policy {linear|lmetric|unified}] [--requests N]"
     echo "  Trace QPS is controlled by sample_trace.py --sample-ratio, not by bench.sh."
     exit 1
 fi
 
+if [ "$MODE" = "elastic" ] && [ "$POLICY_SET" = "false" ]; then
+    POLICY="unified"
+fi
+
 OUTDIR="$PROJECT_DIR/outputs/$TAG"
 if [ -d "$OUTDIR" ] && [ -f "$OUTDIR/metrics.jsonl" ]; then
     echo "[ERROR] Output directory $OUTDIR already exists with data. Use a different --tag."

scripts/cache_aware_proxy.py

@@ -14,6 +14,7 @@ Routing policies (--policy):
 
 import argparse
 import asyncio
+import json
 import os
 import time as _time
 import urllib.parse
@@ -45,6 +46,10 @@ class Settings:
     prefill_throughput: float = 7000.0     # tokens/s per GPU (measured on H20)
     rdma_overhead_s: float = 0.1           # RDMA PUSH overhead (~10-50ms measured)
     cache_capacity_blocks: int = 200000    # per-instance LRU cap on shadow cached_blocks
+    heavy_threshold: int = 20000
+    overload_factor: float = 2.0
+    max_offload_inflight: int = 4
+    cache_gate_ratio: float = 0.0
 
 
 SETTINGS = Settings()
@@ -98,7 +103,7 @@ def _p_offload_penalty(inst: InstanceState) -> int:
     """Penalty for PD-sep mode routing (legacy)."""
     if inst.active_p_offloads <= 0:
         return 0
-    return inst.active_p_offloads * 20000
+    return inst.active_p_offloads * SETTINGS.heavy_threshold
 
 
 def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
@@ -118,7 +123,7 @@ def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
         idx = affinity[session_id]
         if idx < len(instances):
             inst = instances[idx]
-            if (inst.ongoing_tokens <= avg_load * 2.0
+            if (inst.ongoing_tokens <= avg_load * SETTINGS.overload_factor
                     and inst.active_p_offloads == 0):
                 return inst, idx
 
@@ -197,6 +202,54 @@ async def _query_bootstrap_hit(
         return None
 
 
+def _extract_output_token_ids_from_sse(
+    buffer: str,
+    chunk: bytes,
+) -> tuple[str, list[int]]:
+    """Extract vLLM streaming token_ids while preserving the raw stream."""
+    buffer += chunk.decode("utf-8", errors="ignore")
+    complete = buffer.endswith("\n") or buffer.endswith("\r")
+    lines = buffer.splitlines()
+    if complete:
+        buffer = ""
+    elif lines:
+        buffer = lines.pop()
+    else:
+        return buffer, []
+
+    output_ids: list[int] = []
+    for line in lines:
+        line = line.strip()
+        if not line.startswith("data:"):
+            continue
+        data = line[5:].strip()
+        if not data or data == "[DONE]":
+            continue
+        try:
+            payload = json.loads(data)
+        except json.JSONDecodeError:
+            continue
+        choices = payload.get("choices", [])
+        for choice in choices:
+            token_ids = choice.get("token_ids")
+            if isinstance(token_ids, list):
+                output_ids.extend(
+                    int(t) for t in token_ids if isinstance(t, int)
+                )
+    return buffer, output_ids
+
+
+def _realized_tokens(
+    prompt_token_ids: list[int] | None,
+    output_token_ids: list[int],
+) -> list[int] | None:
+    if prompt_token_ids is None:
+        return None
+    if not output_token_ids:
+        return prompt_token_ids
+    return prompt_token_ids + output_token_ids
+
+
 global_args = None
 combined_instances: list[InstanceState] = []
 prefill_instances: list[InstanceState] = []
@@ -367,6 +420,54 @@ async def _handle(request: Request, api: str):
                                        input_length, session_id, headers)
 
 
+async def _handle_local_request(api, req_data, headers, token_ids, input_length,
+                                chosen: InstanceState, estimated_new: int,
+                                breakdown: dict):
+    breakdown.setdefault("route_class", "LOCAL")
+    breakdown.setdefault("routed_to", chosen.url)
+    chosen.ongoing_tokens += input_length
+    chosen.pending_prefill_tokens += estimated_new
+    chosen.num_requests += 1
+
+    async def generate():
+        prefill_done = False
+        sse_buffer = ""
+        output_token_ids: list[int] = []
+        try:
+            for attempt in range(MAX_STREAM_RETRIES):
+                try:
+                    async with chosen.client.stream("POST", api, json=req_data, headers=headers) as resp:
+                        resp.raise_for_status()
+                        async for chunk in resp.aiter_bytes():
+                            sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
+                                sse_buffer, chunk)
+                            output_token_ids.extend(new_output_ids)
+                            if not prefill_done:
+                                chosen.pending_prefill_tokens -= estimated_new
+                                chosen.ongoing_decode_tokens += input_length
+                                breakdown["t_first_token"] = _time.monotonic()
+                                prefill_done = True
+                            yield chunk
+                    chosen.record_prefix(
+                        _realized_tokens(token_ids, output_token_ids))
+                    break
+                except (httpx.ConnectError, httpx.RemoteProtocolError):
+                    if prefill_done or attempt >= MAX_STREAM_RETRIES - 1:
+                        raise
+                    await asyncio.sleep(RETRY_DELAY_S)
+        finally:
+            if not prefill_done:
+                chosen.pending_prefill_tokens -= estimated_new
+            else:
+                chosen.ongoing_decode_tokens -= input_length
+            chosen.ongoing_tokens -= input_length
+            chosen.num_requests -= 1
+            breakdown["t_done"] = _time.monotonic()
+            _breakdown_log.append(breakdown)
+
+    return StreamingResponse(generate(), media_type="text/event-stream")
+
+
 async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers):
     """Unified routing: pick the instance with lowest expected latency.
 
@@ -375,13 +476,57 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
     where prefill_time depends on whether the instance has cache (local),
     can receive cache via PUSH (remote), or must do cold prefill.
     """
+    policy = getattr(global_args, 'policy', 'linear')
     offload_enabled = getattr(global_args, 'offload', False) and len(combined_instances) >= 2
     throughput = SETTINGS.prefill_throughput
 
+    if policy in ("linear", "lmetric"):
+        if policy == "lmetric":
+            chosen, best_idx = pick_instance_lmetric(
+                combined_instances, token_ids, session_id, input_length,
+                session_affinity_combined)
+        else:
+            chosen, best_idx = pick_instance(
+                combined_instances, token_ids, session_id, input_length,
+                session_affinity_combined)
+        cache_hit = chosen.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,
+            "cache_hit": cache_hit,
+            "estimated_new_tokens": estimated_new,
+            "t_proxy_recv": _time.monotonic(),
+            "policy": policy,
+            "route_class": "LOCAL",
+            "routed_to": chosen.url,
+        }
+        if session_id and policy == "lmetric":
+            # LMetric is intentionally per-request; record last target only for
+            # stats/debugging, not for future decisions.
+            session_affinity_combined[session_id] = best_idx
+        return await _handle_local_request(
+            api, req_data, headers, token_ids, input_length,
+            chosen, estimated_new, breakdown)
+
     # Compute cache hits for all instances
     cache_hits = [inst.estimate_cache_hit(token_ids) for inst in combined_instances]
     best_cache_idx = max(range(len(combined_instances)), key=lambda i: cache_hits[i])
     best_cache_hit = cache_hits[best_cache_idx]
+    def _current_offloads() -> int:
+        return sum(i.active_p_offloads for i in combined_instances)
+
+    def _push_allowed(cache_hit: int) -> bool:
+        if _current_offloads() >= SETTINGS.max_offload_inflight:
+            return False
+        push_new = max(0, input_length - cache_hit)
+        if push_new < SETTINGS.heavy_threshold:
+            return False
+        if SETTINGS.cache_gate_ratio > 0:
+            cache_ratio = cache_hit / max(input_length, 1)
+            if cache_ratio < SETTINGS.cache_gate_ratio:
+                return False
+        return True
 
     def _instance_cost(i: int) -> tuple[float, bool]:
         """Expected latency if this request goes to instance i."""
@@ -391,7 +536,8 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
         local_new = max(0, input_length - local_hit)
         local_cost = queue + local_new / throughput
 
-        if offload_enabled and best_cache_hit > 0 and i != best_cache_idx and local_hit < best_cache_hit:
+        if (offload_enabled and best_cache_hit > 0 and _push_allowed(best_cache_hit)
+                and i != best_cache_idx and local_hit < best_cache_hit):
             push_new = max(0, input_length - best_cache_hit)
             push_cost = queue + push_new / throughput + SETTINGS.rdma_overhead_s
             if push_cost < local_cost:
@@ -406,7 +552,7 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
         all_costs = [_instance_cost(i) for i in range(len(combined_instances))]
         global_best_cost = min(c for c, _ in all_costs)
         # Use affinity if it's within 2x of the best option
-        if affinity_cost <= global_best_cost * 2.0:
+        if affinity_cost <= global_best_cost * SETTINGS.overload_factor:
             best_idx = affinity_idx
             best_cost = affinity_cost
             best_needs_push = affinity_push
@@ -428,6 +574,7 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
         "cache_hit": cache_hit,
         "estimated_new_tokens": estimated_new,
         "t_proxy_recv": _time.monotonic(),
+        "policy": policy,
         "chosen_cost": round(best_cost, 2),
     }
 
@@ -451,6 +598,23 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
         # If real hit > 0, proceed with offload
         if push_cache_hit > 0:
             push_new = max(0, input_length - push_cache_hit)
+            cache_ratio = push_cache_hit / max(input_length, 1)
+
+            if _current_offloads() >= SETTINGS.max_offload_inflight:
+                breakdown["push_downgraded"] = "cap_reached"
+                return await _handle_local_request(
+                    api, req_data, headers, token_ids, input_length,
+                    chosen, estimated_new, breakdown)
+            if push_new < SETTINGS.heavy_threshold:
+                breakdown["push_downgraded"] = "below_heavy_threshold"
+                return await _handle_local_request(
+                    api, req_data, headers, token_ids, input_length,
+                    chosen, estimated_new, breakdown)
+            if SETTINGS.cache_gate_ratio > 0 and cache_ratio < SETTINGS.cache_gate_ratio:
+                breakdown["push_downgraded"] = "cache_gate"
+                return await _handle_local_request(
+                    api, req_data, headers, token_ids, input_length,
+                    chosen, estimated_new, breakdown)
 
             offload_mode = getattr(global_args, 'offload_mode', 'cached_prefill')
             breakdown["c_inst"] = c_inst.url
@@ -482,42 +646,9 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
     # LOCAL path (also handles downgraded PUSH)
     breakdown["route_class"] = "LOCAL"
     breakdown["routed_to"] = chosen.url
-
-    chosen.ongoing_tokens += input_length
-    chosen.pending_prefill_tokens += estimated_new
-    chosen.num_requests += 1
-
-    async def generate():
-        prefill_done = False
-        try:
-            for attempt in range(MAX_STREAM_RETRIES):
-                try:
-                    async with chosen.client.stream("POST", api, json=req_data, headers=headers) as resp:
-                        resp.raise_for_status()
-                        async for chunk in resp.aiter_bytes():
-                            if not prefill_done:
-                                chosen.pending_prefill_tokens -= estimated_new
-                                chosen.ongoing_decode_tokens += input_length
-                                breakdown["t_first_token"] = _time.monotonic()
-                                prefill_done = True
-                            yield chunk
-                    chosen.record_prefix(token_ids)
-                    break
-                except (httpx.ConnectError, httpx.RemoteProtocolError):
-                    if prefill_done or attempt >= MAX_STREAM_RETRIES - 1:
-                        raise
-                    await asyncio.sleep(RETRY_DELAY_S)
-        finally:
-            if not prefill_done:
-                chosen.pending_prefill_tokens -= estimated_new
-            else:
-                chosen.ongoing_decode_tokens -= input_length
-            chosen.ongoing_tokens -= input_length
-            chosen.num_requests -= 1
-            breakdown["t_done"] = _time.monotonic()
-            _breakdown_log.append(breakdown)
-
-    return StreamingResponse(generate(), media_type="text/event-stream")
+    return await _handle_local_request(
+        api, req_data, headers, token_ids, input_length,
+        chosen, estimated_new, breakdown)
 
 
 PREFILL_TIMEOUT_S = 120  # max seconds to wait for P-instance prefill
@@ -544,6 +675,7 @@ async def _handle_cached_prefill_offload(api, req_data, headers, token_ids,
     }
     prefill_data["stream"] = False
     prefill_data["max_tokens"] = 1
+    prefill_data["min_tokens"] = 1
     prefill_data.pop("max_completion_tokens", None)
     prefill_data.pop("stream_options", None)
     p_headers = {**headers, "X-data-parallel-rank": "0"}
@@ -591,16 +723,21 @@ async def _handle_cached_prefill_offload(api, req_data, headers, token_ids,
 
     async def generate():
         first_token = True
+        sse_buffer = ""
+        output_token_ids: list[int] = []
         try:
             async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
                 resp.raise_for_status()
                 async for chunk in resp.aiter_bytes():
+                    sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
+                        sse_buffer, chunk)
+                    output_token_ids.extend(new_output_ids)
                     if first_token:
                         d_inst.ongoing_decode_tokens += input_length
                         breakdown["t_first_token"] = _time.monotonic()
                         first_token = False
                     yield chunk
-            d_inst.record_prefix(token_ids)
+            d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
         finally:
             if not first_token:
                 d_inst.ongoing_decode_tokens -= input_length
@@ -641,17 +778,22 @@ async def _handle_direct_read_offload(api, req_data, headers, token_ids,
 
     async def generate():
         first_token = True
+        sse_buffer = ""
+        output_token_ids: list[int] = []
         try:
             async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
                 resp.raise_for_status()
                 async for chunk in resp.aiter_bytes():
+                    sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
+                        sse_buffer, chunk)
+                    output_token_ids.extend(new_output_ids)
                     if first_token:
                         d_inst.pending_prefill_tokens -= estimated_new
                         d_inst.ongoing_decode_tokens += input_length
                         breakdown["t_first_token"] = _time.monotonic()
                         first_token = False
                     yield chunk
-            d_inst.record_prefix(token_ids)
+            d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
         finally:
             if first_token:
                 d_inst.pending_prefill_tokens -= estimated_new
@@ -688,6 +830,7 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
     }
     prefill_data["stream"] = False
     prefill_data["max_tokens"] = 1
+    prefill_data["min_tokens"] = 1
     prefill_data.pop("max_completion_tokens", None)
     prefill_data.pop("stream_options", None)
     p_headers = {**headers, "X-data-parallel-rank": "0"}
@@ -726,14 +869,20 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
 
     async def generate():
         first_token = True
+        sse_buffer = ""
+        output_token_ids: list[int] = []
         try:
             async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
                 resp.raise_for_status()
                 async for chunk in resp.aiter_bytes():
+                    sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
+                        sse_buffer, chunk)
+                    output_token_ids.extend(new_output_ids)
                     if first_token:
                         breakdown["t_first_token"] = _time.monotonic()
                         first_token = False
                     yield chunk
+            d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
         finally:
             breakdown["t_done"] = _time.monotonic()
             d_inst.ongoing_tokens -= input_length
@@ -779,8 +928,9 @@ def parse_args():
                    help="Enable Mooncake KV offload for HEAVY requests (requires kv_both instances)")
     p.add_argument("--bootstrap-ports", type=str, default="",
                    help="Comma-separated bootstrap ports for combined instances (for offload mode)")
-    p.add_argument("--policy", type=str, default="linear", choices=["linear", "lmetric"],
-                   help="Routing policy: linear (default) or lmetric (P_tokens × BS, OSDI'26)")
+    p.add_argument("--policy", type=str, default="linear",
+                   choices=["linear", "lmetric", "unified"],
+                   help="Routing policy: linear, lmetric (P_tokens × BS), or unified cost model")
     p.add_argument("--overload-factor", type=float, default=2.0,
                    help="Break session affinity when instance load > factor * avg")
     p.add_argument("--max-offload-inflight", type=int, default=4,
@@ -789,7 +939,7 @@ def parse_args():
                    choices=["direct_read", "cached_prefill"],
                    help="direct_read: D pulls KV from C (PUSH). "
                         "cached_prefill: C prefills then D decodes (PD-sep style).")
-    p.add_argument("--cache-gate-ratio", type=float, default=0.3,
+    p.add_argument("--cache-gate-ratio", type=float, default=0.0,
                    help="Min cache_hit/input ratio to allow offload "
                         "(0.0 disables gate, 1.0 disables offload entirely)")
     args = p.parse_args()
@@ -809,6 +959,10 @@ def parse_args():
 
 if __name__ == "__main__":
     global_args = parse_args()
+    SETTINGS.heavy_threshold = global_args.heavy_threshold
+    SETTINGS.overload_factor = global_args.overload_factor
+    SETTINGS.max_offload_inflight = global_args.max_offload_inflight
+    SETTINGS.cache_gate_ratio = global_args.cache_gate_ratio
     print("SETTINGS: throughput=%.0f rdma_overhead=%.2f offload=%s" % (
         SETTINGS.prefill_throughput, SETTINGS.rdma_overhead_s,
         getattr(global_args, 'offload', False)))

scripts/launch_elastic_p2p.sh

@@ -90,6 +90,7 @@ $PYTHON "$PROJECT_DIR/scripts/cache_aware_proxy.py" \
     --combined $combined_args \
     --bootstrap-ports "$bootstrap_ports" \
     --offload \
+    --policy unified \
     --heavy-threshold $HEAVY_THRESHOLD \
     --port $PROXY_PORT &
 sleep 5

third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -249,6 +249,10 @@ class MooncakeConnector(KVConnectorBase_V1):
         assert self.connector_worker is not None
         return self.connector_worker.get_finished()
 
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
     def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         assert self.connector_worker is not None
         assert isinstance(self._connector_metadata, MooncakeConnectorMetadata)
@@ -512,10 +516,13 @@ class MooncakeConnectorScheduler:
 
         # TODO: check whether block_ids actually ever be 0. If not we could
         # remove the conditional below
-        delay_free_blocks = len(block_ids) > 0
+        block_size = self.vllm_config.cache_config.block_size
+        prompt_blocks = (request.num_prompt_tokens + block_size - 1) // block_size
+        send_block_ids = block_ids[:prompt_blocks]
+        delay_free_blocks = len(send_block_ids) > 0
 
         if delay_free_blocks:
-            self._reqs_need_send[request.request_id] = (request, block_ids)
+            self._reqs_need_send[request.request_id] = (request, send_block_ids)
 
         return delay_free_blocks, None
 
@@ -620,6 +627,7 @@ class MooncakeConnectorWorker:
 
         self.finished_sending_reqs: set[ReqId] = set()
         self.finished_recving_reqs: set[ReqId] = set()
+        self.failed_recving_block_ids: set[int] = set()
 
         self.block_size = vllm_config.cache_config.block_size
         self.model_config = vllm_config.model_config
@@ -1063,6 +1071,11 @@ class MooncakeConnectorWorker:
         self.finished_recving_reqs = set()
         return finished_recving_reqs
 
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        failed = self.failed_recving_block_ids
+        self.failed_recving_block_ids = set()
+        return failed
+
     async def fetch_finished_sending_reqs(self) -> set[ReqId]:
         finished_sending_reqs = self.finished_sending_reqs
         self.finished_sending_reqs = set()
@@ -1176,6 +1189,10 @@ class MooncakeConnectorWorker:
             logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
         except Exception as e:
             logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
+            for req_id in req_ids:
+                pull_meta = pull_metas[req_id]
+                self.failed_recving_block_ids.update(pull_meta.local_block_ids)
+                self.finished_recving_reqs.add(pull_meta.d_req_id)
             return
 
     def process_pulling_result(
@@ -1201,6 +1218,12 @@ class MooncakeConnectorWorker:
                 response.err_reqs,
                 response.err_msg,
             )
+            for req_id in response.err_reqs:
+                pull_meta = pull_metas.get(req_id)
+                if pull_meta is None:
+                    continue
+                self.failed_recving_block_ids.update(pull_meta.local_block_ids)
+                self.finished_recving_reqs.add(pull_meta.d_req_id)
 
     async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
         url = remote_bootstrap_addr + "/query"
@@ -1322,11 +1345,13 @@ class MooncakeConnectorWorker:
                 logger.info("direct_push %s: %d blocks pushed from C", req_id, matched)
             else:
                 logger.debug("direct_push %s: %d matched, pushed=%s", req_id, matched, pushed)
+                self.failed_recving_block_ids.update(local_block_ids)
 
             self.finished_recving_reqs.add(req_id)
 
         except Exception as e:
             logger.error("direct_push %s failed: %s", req_id, e)
+            self.failed_recving_block_ids.update(pm.local_block_ids)
             self.finished_recving_reqs.add(req_id)
 
     async def _start_load_kv(