"""Unified cache-aware + token-level load-balanced global scheduler.

Supports two modes:
  --combined URL [URL ...]: PD co-located instances (normal vLLM, no KV transfer)
  --prefill URL BP --decode URL: PD disaggregated instances (Mooncake KV transfer)

Routing policies (--policy):
  linear (default): score = ongoing_tokens - ALPHA * cache_hit_tokens
  lmetric: score = P_tokens * BS  (LMetric, OSDI'26)
           P_tokens = pending_prefill_tokens + new_uncached_tokens
           BS = num_requests (waiting + running)
  Session affinity: multi-turn sessions stick to same instance (all policies).
"""

import argparse
import asyncio
import json
import os
import time as _time
import urllib.parse
import uuid
from collections import OrderedDict, deque
from contextlib import asynccontextmanager
from dataclasses import dataclass

import httpx

MAX_STREAM_RETRIES = 3
RETRY_DELAY_S = 0.5
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse

BLOCK_SIZE = 512
CACHE_HIT_ALPHA = 1.0


@dataclass
class Settings:
    """Runtime-tunable knobs. Populated from argparse in __main__.

    All routing/offload code reads from the SETTINGS singleton so that
    CLI overrides survive even when the module is imported as a library
    (e.g. by tests/) and __main__ does not run.
    """
    prefill_throughput: float = 7000.0     # tokens/s per GPU (measured on H20)
    rdma_overhead_s: float = 0.1           # legacy floor; v2 uses estimate_transfer_cost
    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
    decode_iteration_s: float = 0.05      # per-request decode iteration cost (H20)

    # --- Patch 6.9: cost-model calibration for unified_v2 ---
    # Throughput when the engine runs in kv_both mode. Lower than the
    # pure-decode 7000 tok/s because kv_both adds always-on overhead
    # (REPORT §3.8 documents ~+16% TPOT vs plain).
    prefill_throughput_kv_both: float = 4000.0
    # Calibrated RDMA transfer cost: base + bandwidth term.
    # Floor from isolated test ≈ 0.3 s (handshake + scheduler step).
    # Bandwidth term reflects realized effective throughput, not
    # theoretical 25 GB/s — production p50 = 1.1 s for ~3 GB ≈ 2.7 GB/s
    # effective on the contended kv_both path. v2 uses this lookup
    # rather than the constant rdma_overhead_s.
    rdma_base_overhead_s: float = 0.3
    rdma_effective_gb_per_s: float = 2.7

    # Qwen3-Coder-30B-A3B (bf16, 48 layers × 4 KV heads × 128 head_dim × 2):
    # 2 × 48 × 4 × 128 × 2 = 98304 bytes per token.
    kv_bytes_per_token: int = 98304

    # --- unified_v2 gating knobs (relaxed in v2.1 after the v1 0.2% trigger rate) ---
    # B2 microbench shows TPOT idx 1.9x already at new_tokens=8k and TTFT
    # idx ~12x; the previous 16k threshold was too conservative and
    # rejected 88.7% of candidates (window_1_results/v2_breakdown).
    pd_sep_min_new_tokens: int = 8000
    pd_sep_min_decodes_protected: int = 1       # any in-flight work on chosen counts
    pd_sep_min_src_cache_tokens: int = 4000     # half a block; was 8000
    pd_sep_min_extra_cache_tokens: int = 2000   # half a block; was 4000
    pd_sep_margin_s: float = 0.2                # require cost gap > 0.2 s before migrating
    # Patch 6.6: per-request KV-xfer wall-clock timeout (proxy side).
    pd_sep_xfer_timeout_s: float = 60.0

    # --- unified_v3 (offload-decode) gating knobs -----------------------
    # v3 differs from v2 in *direction*: prefill stays on the session-
    # affinity host (which holds the prefix cache); decode is migrated to
    # a less-loaded target. KV transfer flows prefill_host → decode_target.
    # The target doesn't need cache — we're shipping the post-prefill KV
    # over anyway. After successful migration the session affinity table
    # rotates to decode_target so the *next* turn lands where the KV now
    # lives.
    v3_min_new_tokens: int = 8000               # same as v2: don't migrate tiny prefills
    v3_min_prefill_decode_busy: int = 1         # prefill_host must have ≥ this many concurrent decode tokens to justify migrating
    v3_target_load_ratio: float = 0.7           # target.num_requests must be < prefill_host.num_requests × this
    v3_min_load_gap: int = 1                    # target.num_requests must also be ≤ prefill_host - this (absolute slack)
    v3_rotate_affinity: bool = True             # after migration, set session affinity to decode_target.
                                                # Empirically False is better — see cache_miss_audit (next turn hits 9.5%
                                                # with rotation vs ~80% without), because delay_free_blocks doesn't
                                                # actually preserve cross-turn KV on decode_target.
    v3_prefer_cache_target: bool = True         # Mechanism B: among low-load candidates, prefer the one
                                                # with the most prefix cache for this prompt — vLLM's connector
                                                # auto-transfers only the missing portion (verified via
                                                # smoke_partial_transfer: cache-rich dst is 77% faster than
                                                # cold dst at 33k tokens, +512 ext).
    # Anti-hotspot: picker scores effective_load = num_requests + (recent
    # migrations received within window). Prevents clustering migrations on
    # one instance in rapid succession (observed in Mech B run: inst_5 became
    # a hotspot via post-rotation tail accumulation).
    v3_recent_mig_window_s: float = 10.0        # sliding window
    v3_recent_mig_weight: float = 1.0           # how many "virtual requests" each
                                                # recent migration counts as

    # P2: real engine-state feed (replaces 30s-stale shadow counters for
    # migration target selection). Empty = disabled (use shadow only).
    engine_state_uri: str = ""                  # file:///dev/shm/... or redis://...
    engine_state_period_ms: int = 50            # router poll period
    es_big_prefill_threshold: int = 16000       # target mid-prefill >= this => avoid (GIL stall)
    es_kv_wall_frac: float = 0.90               # target KV usage >= this => avoid (capacity wall)

    # Direction B knob: LMetric fallback adds decode-token penalty to score.
    # score = (pending_prefill + new + lmetric_decode_weight * ongoing_decode_tok) * num_req
    # Empirical iter-time slope on H100 + Qwen3-30B-A3B: each decode token in
    # batch costs ~0.01 prefill-token-equivalent in scheduler time, so 0.01 is
    # a reasonable starting weight. Set 0 to disable (original behavior).
    lmetric_decode_weight: float = 0.0

    # --- KV connector selection (governs PD-sep handshake) -------------
    # "mooncake": pre-baked kv_transfer_params (bootstrap_addr+engine_id+transfer_id).
    #             Requires --bootstrap-ports and vLLMs launched with MooncakeConnector.
    # "nixl"    : response-forward handshake. src returns kv_transfer_params via
    #             response body, proxy forwards to dst. Nixl auto-selects transport
    #             via UCX (CUDA IPC / NVLink on intra-node, RDMA across nodes).
    connector_type: str = "mooncake"


SETTINGS = Settings()


def estimate_transfer_cost(transfer_bytes: int) -> float:
    """Calibrated RDMA transfer cost as a function of bytes.

    Replaces the legacy constant rdma_overhead_s. Calibration sources:
    - Floor: isolated-test ~0.3 s for a few-block PUSH (scripts/test_direct_read.py)
    - Bandwidth term: outputs/contention_16s_elastic/breakdown.json shows
      decode_sent->first_token p50 = 1.1 s for ~3 GB transfers, giving
      ~2.7 GB/s effective on the contended kv_both path.

    The p90 in that same run is 6.7 s (D-side block reservation +
    scheduler step delays). v2's cost model uses the *median* — being
    too pessimistic would suppress all PD-sep triggers. The risk of
    underestimation is mitigated by the pd_sep_margin_s safety factor.
    """
    base = SETTINGS.rdma_base_overhead_s
    bw_term = transfer_bytes / (SETTINGS.rdma_effective_gb_per_s * 1024 ** 3)
    return base + bw_term


def estimate_same_worker_interference_s(
    new_tokens: int,
    num_decodes: int,
) -> float:
    """Estimated additional latency on `num_decodes` co-located decodes
    when a `new_tokens`-token prefill runs on the same worker.

    Derived from B2 microbench (analysis/characterization/window_1_results.md):
    same-worker prefill of size N steals decode capacity for the
    prefill's duration. The penalty factor is the fraction of decode
    steps stolen during the prefill window.

    For new_tokens < 4k:   ~0.2  (chunked prefill leaves room)
    For new_tokens 16k:    ~0.5  (mid-regime, B2 TPOT idx 3.4×)
    For new_tokens 32k:    ~0.8  (B2 peak TPOT idx 7.9×)
    For new_tokens > 32k:  ~0.95 (B2 TTFT regime — decodes are nearly fully blocked)

    The cost in seconds is roughly: prefill_duration × penalty × n_decodes,
    because each affected decode loses ~penalty fraction of its capacity
    during the prefill window.
    """
    if num_decodes <= 0:
        return 0.0
    prefill_dur_s = new_tokens / SETTINGS.prefill_throughput_kv_both
    if new_tokens < 4000:
        penalty = 0.2
    elif new_tokens < 16000:
        penalty = 0.5
    elif new_tokens < 32000:
        penalty = 0.8
    else:
        penalty = 0.95
    return prefill_dur_s * penalty * num_decodes


class InstanceState:
    def __init__(self, url: str, bootstrap_port: int | None = None):
        self.url = url
        self.bootstrap_port = bootstrap_port
        self.client = httpx.AsyncClient(
            timeout=None, base_url=url,
            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.pending_prefill_tokens = 0  # tokens for requests still in prefill
        self.num_requests = 0  # total in-flight requests (waiting + running)
        self.active_p_offloads = 0  # number of HEAVY prefills this instance is doing for others
        self.engine_id: dict[int, str] = {}
        self.dp_size = 1
        # OrderedDict acts as an LRU keyed by block hash; value is unused.
        self.cached_blocks: OrderedDict[int, None] = OrderedDict()
        # v3 anti-hotspot: timestamps (monotonic) when this instance was picked
        # as a v3 migration target. Used to compute effective_load = num_req +
        # recent-migration count over a sliding window, preventing back-to-back
        # decisions from clustering on the same dst.
        self.recent_mig_targeted_at: deque[float] = deque(maxlen=64)
        # P2: latest real engine state (from the engine-state feed), or None
        # when the feed is disabled/stale. Set by _engine_state_poll_loop.
        self.real_state: dict | None = None

    # ---- effective-load accessors (P2): prefer REAL engine state when the
    # feed is fresh, else the proxy shadow counter. We take max(shadow, real)
    # for load so we never under-count: REAL fixes the 30s-stale under-count,
    # while the shadow's atomic pre-await reservation still covers the
    # in-flight window (preserving the RaceFix against concurrent picks).
    def eff_num_requests(self) -> float:
        rs = self.real_state
        if rs is not None:
            return max(self.num_requests,
                       rs.get("num_running", 0) + rs.get("num_waiting", 0))
        return self.num_requests

    def eff_pending_prefill(self) -> float:
        rs = self.real_state
        if rs is not None:
            return max(self.pending_prefill_tokens,
                       rs.get("pending_prefill_tokens", 0))
        return self.pending_prefill_tokens

    def eff_ongoing_decode(self) -> float:
        rs = self.real_state
        if rs is not None:
            return max(self.ongoing_decode_tokens,
                       rs.get("ongoing_decode_tokens", 0))
        return self.ongoing_decode_tokens

    def eff_ongoing_tokens(self) -> float:
        rs = self.real_state
        if rs is not None:
            return max(self.ongoing_tokens,
                       rs.get("pending_prefill_tokens", 0)
                       + rs.get("ongoing_decode_tokens", 0))
        return self.ongoing_tokens

    def r_max_prefill_remaining(self) -> int:
        rs = self.real_state
        return int(rs.get("max_prefill_remaining", 0)) if rs is not None else 0

    def r_kv_used_frac(self) -> float:
        rs = self.real_state
        if rs is not None:
            return float(rs.get("gpu_kv_used_frac", 0.0) or 0.0)
        return 0.0

    def estimate_cache_hit(self, token_ids: list[int] | None) -> int:
        if not token_ids or len(token_ids) < BLOCK_SIZE:
            return 0
        hit = 0
        for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
            bh = hash(tuple(token_ids[i:i + BLOCK_SIZE]))
            if bh in self.cached_blocks:
                self.cached_blocks.move_to_end(bh)  # LRU touch on hit
                hit += BLOCK_SIZE
            else:
                break
        return hit

    def record_prefix(self, token_ids: list[int] | None):
        if not token_ids:
            return
        for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
            bh = hash(tuple(token_ids[i:i + BLOCK_SIZE]))
            if bh in self.cached_blocks:
                self.cached_blocks.move_to_end(bh)
            else:
                self.cached_blocks[bh] = None
                if len(self.cached_blocks) > SETTINGS.cache_capacity_blocks:
                    self.cached_blocks.popitem(last=False)


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 * SETTINGS.heavy_threshold


def snapshot_workers(
    instances: list[InstanceState],
    token_ids: list[int] | None = None,
    input_length: int = 0,
) -> list[dict]:
    """Per-worker state at route-decision time.

    All routing-relevant counters plus the score each policy would
    have produced for `input_length` if it were dispatched now. Cheap
    enough to call on every request; B3 hot-spot analysis depends on
    this being captured per decision.
    """
    snap: list[dict] = []
    for i, inst in enumerate(instances):
        cache_hit = inst.estimate_cache_hit(token_ids) if token_ids else 0
        new_prefill = max(0, input_length - cache_hit)
        snap.append({
            "idx": i,
            "url": inst.url,
            "ongoing_tokens": inst.ongoing_tokens,
            "ongoing_decode_tokens": inst.ongoing_decode_tokens,
            "pending_prefill_tokens": inst.pending_prefill_tokens,
            "num_requests": inst.num_requests,
            "active_p_offloads": inst.active_p_offloads,
            "cached_blocks": len(inst.cached_blocks),
            "cache_hit": cache_hit,
            "new_prefill": new_prefill,
            # scores reflect the ACTUAL decision basis (eff = real-or-shadow).
            "score_linear": (inst.eff_ongoing_tokens()
                              + _p_offload_penalty(inst)
                              - CACHE_HIT_ALPHA * cache_hit),
            "score_lmetric": (inst.eff_pending_prefill() + new_prefill)
                              * inst.eff_num_requests(),
            # P2: real-state fields when the feed is fresh (None otherwise).
            "real_num_requests": (inst.eff_num_requests()
                                  if inst.real_state is not None else None),
            "real_max_prefill_remaining": (inst.r_max_prefill_remaining()
                                           if inst.real_state is not None else None),
            "real_kv_used_frac": (inst.r_kv_used_frac()
                                  if inst.real_state is not None else None),
        })
    return snap


def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
                  session_id: str | None, input_length: int,
                  affinity: dict[str, int]) -> tuple[InstanceState, int]:
    """Session-sticky with load-aware override.

    Turn 2+: use session affinity UNLESS pinned instance is overloaded
    or busy with P-role offloads, in which case pick least-loaded.
    Turn 1: pick instance with best score (load + cache combined).
    Instances doing P-role offloads get a large penalty to steer
    WARM/MEDIUM traffic away.
    """
    avg_load = max(sum(i.eff_ongoing_tokens() for i in instances) / len(instances), 1.0)

    if session_id and session_id in affinity:
        idx = affinity[session_id]
        if idx < len(instances):
            inst = instances[idx]
            if (inst.eff_ongoing_tokens() <= avg_load * SETTINGS.overload_factor
                    and inst.active_p_offloads == 0):
                return inst, idx

    best_idx, best_score = 0, float("inf")
    for i, inst in enumerate(instances):
        cache_hit = inst.estimate_cache_hit(token_ids)
        score = (inst.eff_ongoing_tokens() + _p_offload_penalty(inst)
                 - CACHE_HIT_ALPHA * cache_hit)
        if score < best_score:
            best_score = score
            best_idx = i

    if session_id:
        affinity[session_id] = best_idx
    return instances[best_idx], best_idx


def pick_instance_load_only(
    instances: list[InstanceState],
    token_ids: list[int] | None,
    session_id: str | None,
    input_length: int,
    affinity: dict[str, int],
) -> tuple[InstanceState, int]:
    """Pure load balancing: pick instance with fewest in-flight requests.

    Ignores cache hits and session affinity. Used as a B3 control to
    isolate the locality contribution of cache-aware policies.
    """
    best_idx = min(range(len(instances)),
                    key=lambda i: instances[i].eff_num_requests())
    return instances[best_idx], best_idx


def pick_instance_sticky(
    instances: list[InstanceState],
    token_ids: list[int] | None,
    session_id: str | None,
    input_length: int,
    affinity: dict[str, int],
) -> tuple[InstanceState, int]:
    """Hard session affinity: once assigned, never break.

    First turn of a session picks the instance with the lowest
    num_requests; subsequent turns always return to the same instance
    regardless of load. Used as a B3 control to isolate the hot-spot
    cost of perfect locality.
    """
    if session_id and session_id in affinity:
        idx = affinity[session_id]
        if idx < len(instances):
            return instances[idx], idx
    best_idx = min(range(len(instances)),
                    key=lambda i: instances[i].eff_num_requests())
    if session_id:
        affinity[session_id] = best_idx
    return instances[best_idx], best_idx


def pick_instance_lmetric(instances: list[InstanceState], token_ids: list[int] | None,
                          session_id: str | None, input_length: int,
                          affinity: dict[str, int]) -> tuple[InstanceState, int]:
    """LMetric routing: score = P_tokens × BS (OSDI'26).

    Pure per-request load-based routing, no session affinity (the
    session_id/affinity args are accepted for signature compatibility
    with pick_instance/pick_instance_unified_hybrid but ignored).
    P = pending_prefill_tokens + (input_length - cache_hit)
    BS = num_requests (current batch size)
    """
    best_idx, best_score = 0, float("inf")
    for i, inst in enumerate(instances):
        cache_hit = inst.estimate_cache_hit(token_ids)
        new_prefill = max(0, input_length - cache_hit)
        p_tokens = inst.eff_pending_prefill() + new_prefill
        bs = inst.eff_num_requests()
        score = p_tokens * bs
        if score < best_score:
            best_score = score
            best_idx = i

    return instances[best_idx], best_idx


_unified_fallback_rr_counter = 0


def pick_instance_unified_hybrid(
    instances: list[InstanceState],
    token_ids: list[int] | None,
    session_id: str | None,
    input_length: int,
    affinity: dict[str, int],
) -> tuple[InstanceState, int, dict]:
    """Hybrid routing: high-cache affinity, else LMetric with tie-breaker.

    Affinity gate (both must hold to stick):
      - affinity instance cache_hit / input_length > 0.5
      - affinity.num_requests <= avg_num_requests * SETTINGS.overload_factor

    Fallback ordering (when affinity not used):
      primary:   score = P_tokens * BS    (LMetric)
      secondary: new_uncached_tokens      (prefer instance with most cache)
      tertiary:  num_requests             (prefer least-loaded)
      quaternary: round-robin             (avoid degenerate inst-0 pinning
                                           when BS=0 across the board)

    Returns (chosen, idx, decision_dict). decision_dict carries the
    review #7 breakdown fields so the caller can merge them verbatim.
    """
    global _unified_fallback_rr_counter
    n = len(instances)
    avg_reqs = max(sum(i.eff_num_requests() for i in instances) / n, 1.0)

    decision: dict = {
        "decision": "lmetric_fallback",
        "affinity_idx": None,
        "chosen_idx": None,
        "affinity_cache_hit": None,
        "affinity_cache_ratio": None,
        "affinity_num_requests": None,
        "avg_num_requests": avg_reqs,
        "fallback_score": None,
        "tie_break_used": False,
    }

    if session_id and session_id in affinity:
        a_idx = affinity[session_id]
        if a_idx < n:
            a_inst = instances[a_idx]
            a_hit = a_inst.estimate_cache_hit(token_ids)
            a_ratio = a_hit / max(input_length, 1)
            decision["affinity_idx"] = a_idx
            decision["affinity_cache_hit"] = a_hit
            decision["affinity_cache_ratio"] = a_ratio
            decision["affinity_num_requests"] = a_inst.eff_num_requests()
            if (a_ratio > 0.5
                    and a_inst.eff_num_requests() <= avg_reqs * SETTINGS.overload_factor):
                decision["decision"] = "affinity"
                decision["chosen_idx"] = a_idx
                return a_inst, a_idx, decision

    # Direction B: extend LMetric with decode-load awareness.
    # Original score = (pending_prefill + new_uncached) * num_requests, which
    # ignores ongoing decode work. A host with 200k decode tokens looks "ideal"
    # (P_tokens=0) but its decode iters are slow due to large batch KV reads.
    #
    # First attempt (BUG): score = (p_tokens + decode_pen) * num_req — when
    # num_req=0 the decode_pen is zeroed out, so idle-but-decoding hosts still
    # look free and accumulate cold prefills (8007 hotspot in A+B v1 run).
    #
    # Fix: max(num_req, 1) so decode_pen contributes on idle hosts too.
    keys: list[tuple[float, int, int, int]] = []
    for i, inst in enumerate(instances):
        cache_hit = inst.estimate_cache_hit(token_ids)
        new_prefill = max(0, input_length - cache_hit)
        # P2: effective (real-or-shadow) load signals.
        p_tokens = inst.eff_pending_prefill() + new_prefill
        decode_pen = SETTINGS.lmetric_decode_weight * inst.eff_ongoing_decode()
        bs = inst.eff_num_requests()
        score = (p_tokens + decode_pen) * max(bs, 1)
        keys.append((score, new_prefill, bs, i))

    best_triple = min(k[:3] for k in keys)
    tied = [k for k in keys if k[:3] == best_triple]
    if len(tied) > 1:
        decision["tie_break_used"] = True
        _unified_fallback_rr_counter += 1
        winner = tied[_unified_fallback_rr_counter % len(tied)]
    else:
        winner = tied[0]
    chosen_idx = winner[3]
    decision["fallback_score"] = winner[0]
    decision["chosen_idx"] = chosen_idx
    return instances[chosen_idx], chosen_idx, decision


def pick_instance_unified_v2(
    instances: list[InstanceState],
    token_ids: list[int] | None,
    session_id: str | None,
    input_length: int,
    affinity: dict[str, int],
) -> tuple[InstanceState, int, dict, tuple[InstanceState, int] | None]:
    """unified_v2 = unified hybrid + selective per-request PD-sep trigger.

    Stage 1 picks `chosen` exactly as `pick_instance_unified_hybrid`.

    Stage 2 asks: is there another instance with materially more cache
    for this request? If yes, would doing prefill on that instance and
    transferring KV to `chosen` for decode be cheaper than just doing
    everything on `chosen`?

    The cost model compares two scenarios in seconds-of-decode-disruption:

      local:  same-worker prefill on chosen of (input - chosen.cache_hit)
              tokens interferes with chosen.num_decodes co-located decodes.

      pd-sep: same-worker prefill on src of (input - src.cache_hit) tokens
              (smaller, because src has more cache) interferes with
              src.num_decodes co-located decodes, plus we pay RDMA
              transfer of src.cache_hit blocks to chosen.

    We migrate only when local cost > pd-sep cost + safety margin AND
    a set of hard gates (size, cache, decodes) are met.

    Returns (chosen, chosen_idx, decision, pd_sep). When pd_sep is None
    the handler should do local routing on `chosen`. When pd_sep is
    (src_inst, src_idx) the handler should do prefill-on-src,
    decode-on-chosen via Mooncake.
    """
    chosen, chosen_idx, decision = pick_instance_unified_hybrid(
        instances, token_ids, session_id, input_length, affinity)

    decision["v2_pd_sep"] = False
    decision["v2_decision"] = "local"
    decision["v2_reason"] = None

    if not token_ids:
        decision["v2_reason"] = "no_token_ids"
        return chosen, chosen_idx, decision, None

    chosen_cache_hit = chosen.estimate_cache_hit(token_ids)
    new_local = max(0, input_length - chosen_cache_hit)

    # Hard gate 1: prefill must be large enough that interference
    # outweighs the fixed RDMA setup cost.
    if new_local < SETTINGS.pd_sep_min_new_tokens:
        decision["v2_reason"] = f"new_local_below_threshold ({new_local} < {SETTINGS.pd_sep_min_new_tokens})"
        return chosen, chosen_idx, decision, None

    # Hard gate 2: chosen must have live decoding work to protect.
    # v2.1 simplification: pure ongoing_decode_tokens check. The previous
    # gate combined num_requests and decode_tokens with AND, but
    # num_requests includes requests still in prefill — adding a prefill
    # to a chosen that has only its own prefill running doesn't disrupt
    # any decode, so skipping makes sense. The right semantic is "skip
    # iff no decode is currently happening on chosen".
    if chosen.ongoing_decode_tokens == 0:
        decision["v2_reason"] = (
            f"chosen_no_active_decode "
            f"(num_req={chosen.num_requests} decode_tok={chosen.ongoing_decode_tokens})"
        )
        return chosen, chosen_idx, decision, None

    # Find best alternative cache source.
    best_src_idx, best_src_hit = -1, 0
    for i, inst in enumerate(instances):
        if i == chosen_idx:
            continue
        h = inst.estimate_cache_hit(token_ids)
        if h > best_src_hit:
            best_src_idx, best_src_hit = i, h

    # Hard gate 3: src must hold meaningful cache.
    if best_src_hit < SETTINGS.pd_sep_min_src_cache_tokens:
        decision["v2_reason"] = f"src_cache_below_threshold ({best_src_hit} < {SETTINGS.pd_sep_min_src_cache_tokens})"
        return chosen, chosen_idx, decision, None

    # Hard gate 4: src must hold materially more cache than chosen.
    if best_src_hit - chosen_cache_hit < SETTINGS.pd_sep_min_extra_cache_tokens:
        decision["v2_reason"] = (
            f"src_not_meaningfully_more_cache "
            f"(src={best_src_hit} chosen={chosen_cache_hit})"
        )
        return chosen, chosen_idx, decision, None

    src = instances[best_src_idx]
    new_src = max(0, input_length - best_src_hit)

    # Cost-benefit in seconds-of-decode-disruption.
    cost_local = estimate_same_worker_interference_s(
        new_local, chosen.num_requests)
    cost_src_interf = estimate_same_worker_interference_s(
        new_src, src.num_requests)
    transfer_bytes = best_src_hit * SETTINGS.kv_bytes_per_token
    cost_xfer = estimate_transfer_cost(transfer_bytes)
    cost_migrate = cost_src_interf + cost_xfer

    decision["v2_chosen_cache_hit"] = chosen_cache_hit
    decision["v2_src_idx"] = best_src_idx
    decision["v2_src_cache_hit"] = best_src_hit
    decision["v2_new_local"] = new_local
    decision["v2_new_src"] = new_src
    decision["v2_cost_local_s"] = cost_local
    decision["v2_cost_src_interf_s"] = cost_src_interf
    decision["v2_cost_xfer_s"] = cost_xfer
    decision["v2_cost_migrate_s"] = cost_migrate

    if cost_local > cost_migrate + SETTINGS.pd_sep_margin_s:
        decision["v2_pd_sep"] = True
        decision["v2_decision"] = "pd_sep"
        decision["v2_reason"] = (
            f"local_cost {cost_local:.2f}s > migrate_cost {cost_migrate:.2f}s "
            f"+ margin {SETTINGS.pd_sep_margin_s:.2f}s"
        )
        return chosen, chosen_idx, decision, (src, best_src_idx)

    decision["v2_reason"] = (
        f"local_cost {cost_local:.2f}s <= migrate_cost {cost_migrate:.2f}s "
        f"+ margin {SETTINGS.pd_sep_margin_s:.2f}s"
    )
    return chosen, chosen_idx, decision, None


def pick_instance_unified_v3(
    instances: list[InstanceState],
    token_ids: list[int] | None,
    session_id: str | None,
    input_length: int,
    affinity: dict[str, int],
) -> tuple[InstanceState, int, dict, tuple[InstanceState, int] | None]:
    """unified_v3 = unified hybrid + selective DECODE migration.

    Direction-reversed from unified_v2:
      - prefill stays on session-affinity host (`prefill_host`) so we keep
        the 93%-intra-session prefix-cache reuse intact.
      - decode is migrated to a lower-load `decode_target` when the
        affinity host is busy with concurrent decodes.
      - KV transfer flows prefill_host → decode_target (the opposite of
        v2's src → chosen).
      - target does NOT need pre-existing cache — we're shipping the
        post-prefill KV over anyway.
      - On successful migration the *caller* must rotate
        `affinity[session_id] = decode_target_idx` so the next turn lands
        where the KV now lives (decode_target retains the blocks after
        completion, since mooncake defaults to delay_free_blocks=True).

    Decision is purely load-based on the target side:
      1. new_local ≥ v3_min_new_tokens   (don't pay RDMA for tiny prefills)
      2. prefill_host.ongoing_decode_tokens ≥ v3_min_prefill_decode_busy
         (the host is actually busy decoding; migration buys decode-bw)
      3. ∃ target with both num_requests < prefill_host.num_requests × ratio
         and num_requests ≤ prefill_host.num_requests − v3_min_load_gap

    Returns (prefill_host, prefill_idx, decision, migrate). When migrate
    is None the request is fully local on prefill_host. When migrate is
    (decode_target_inst, decode_target_idx), the handler should run
    prefill on prefill_host and ship KV to decode_target for decode.
    """
    prefill_host, prefill_idx, decision = pick_instance_unified_hybrid(
        instances, token_ids, session_id, input_length, affinity)

    decision["v3_migrate"] = False
    decision["v3_decision"] = "local"
    decision["v3_reason"] = None

    if not token_ids:
        decision["v3_reason"] = "no_token_ids"
        return prefill_host, prefill_idx, decision, None

    prefill_cache_hit = prefill_host.estimate_cache_hit(token_ids)
    new_local = max(0, input_length - prefill_cache_hit)
    decision["v3_prefill_cache_hit"] = prefill_cache_hit
    decision["v3_new_local"] = new_local

    # Gate 1: prefill must be large enough to amortise RDMA setup.
    if new_local < SETTINGS.v3_min_new_tokens:
        decision["v3_reason"] = (
            f"new_local_below_threshold ({new_local} < {SETTINGS.v3_min_new_tokens})"
        )
        return prefill_host, prefill_idx, decision, None

    # Gate 2: affinity host must be busy with concurrent decodes — that's
    # what migrating decode-traffic-away buys us. If the host is idle
    # there's no point.
    if prefill_host.eff_ongoing_decode() < SETTINGS.v3_min_prefill_decode_busy:
        decision["v3_reason"] = (
            f"prefill_host_not_busy "
            f"(ongoing_decode_tokens={prefill_host.eff_ongoing_decode()} < "
            f"{SETTINGS.v3_min_prefill_decode_busy})"
        )
        return prefill_host, prefill_idx, decision, None

    # Gate 3: pick the lowest-effective-load target. effective_load adds a
    # penalty for recent migrations the instance has received (anti-hotspot).
    now_mono = _time.monotonic()
    cutoff = now_mono - SETTINGS.v3_recent_mig_window_s

    def _real_load(inst):
        # P2: effective (real-or-shadow) request load; see eff_num_requests.
        return inst.eff_num_requests()

    def effective_load(inst):
        # Drop expired entries lazily.
        while inst.recent_mig_targeted_at and inst.recent_mig_targeted_at[0] < cutoff:
            inst.recent_mig_targeted_at.popleft()
        recent = len(inst.recent_mig_targeted_at)
        return _real_load(inst) + recent * SETTINGS.v3_recent_mig_weight

    ph_load = _real_load(prefill_host)
    threshold_loaded = max(1,
        int(ph_load * SETTINGS.v3_target_load_ratio))
    candidates = [
        (i, inst) for i, inst in enumerate(instances)
        if i != prefill_idx
        and effective_load(inst) < threshold_loaded
        and effective_load(inst) <= ph_load - SETTINGS.v3_min_load_gap
    ]
    if not candidates:
        decision["v3_reason"] = (
            f"no_low_load_target "
            f"(prefill_host.num_req={prefill_host.num_requests} "
            f"threshold={threshold_loaded} "
            f"eff_loads=[{','.join(f'{int(effective_load(i))}' for i in instances)}])"
        )
        return prefill_host, prefill_idx, decision, None

    # Mechanism B (v3_prefer_cache_target=True): rank candidates first by
    # cache_hit DESC (more cache = less KV to transfer), then by effective_load
    # (which includes recent-migration penalty), then by ongoing_tokens.
    if SETTINGS.v3_prefer_cache_target:
        def _tgt_key(x):
            # P2: avoid a target that is mid-large-prefill (holds the GIL,
            # stalls the mooncake receiver_loop = the ~45% control-plane
            # residual layer-wise can't fix) or near the KV capacity wall,
            # before ranking by cache-richness and real load.
            inst = x[1]
            ch = inst.estimate_cache_hit(token_ids)
            rs = getattr(inst, "real_state", None)
            stalls = near_wall = 0
            if rs is not None:
                if int(rs.get("max_prefill_remaining", 0)) >= SETTINGS.es_big_prefill_threshold:
                    stalls = 1
                f = rs.get("gpu_kv_used_frac", 0.0) or 0.0
                if float(f) >= SETTINGS.es_kv_wall_frac:
                    near_wall = 1
            return (stalls, near_wall, -ch, effective_load(inst), inst.ongoing_tokens)
        decode_target_idx, decode_target = min(candidates, key=_tgt_key)
    else:
        decode_target_idx, decode_target = min(
            candidates, key=lambda x: (effective_load(x[1]), x[1].ongoing_tokens))

    target_cache_hit = decode_target.estimate_cache_hit(token_ids)
    target_recent_received = len(decode_target.recent_mig_targeted_at)
    # Record this decision for the anti-hotspot accounting.
    decode_target.recent_mig_targeted_at.append(now_mono)

    decision["v3_migrate"] = True
    decision["v3_decision"] = "migrate_decode"
    decision["v3_src_idx"] = prefill_idx
    decision["v3_target_idx"] = decode_target_idx
    decision["v3_target_num_req"] = decode_target.num_requests
    decision["v3_target_cache_hit"] = target_cache_hit
    decision["v3_target_recent_received"] = target_recent_received
    decision["v3_prefill_num_req"] = prefill_host.num_requests
    # Snapshot of src state at the moment of decision (for postmortem).
    decision["v3_src_state"] = {
        "num_requests": prefill_host.num_requests,
        "ongoing_tokens": prefill_host.ongoing_tokens,
        "ongoing_decode_tokens": prefill_host.ongoing_decode_tokens,
        "pending_prefill_tokens": prefill_host.pending_prefill_tokens,
    }
    decision["v3_target_state"] = {
        "num_requests": decode_target.num_requests,
        "ongoing_tokens": decode_target.ongoing_tokens,
        "ongoing_decode_tokens": decode_target.ongoing_decode_tokens,
        "pending_prefill_tokens": decode_target.pending_prefill_tokens,
        "cache_hit_estimate": target_cache_hit,
        "recent_mig_received_in_window": target_recent_received,
    }
    decision["v3_reason"] = (
        f"prefill_host.num_req={prefill_host.num_requests} busy; "
        f"target.num_req={decode_target.num_requests} cache_hit={target_cache_hit} "
        f"recent_received={target_recent_received}, "
        f"transferring KV after prefill"
    )
    return prefill_host, prefill_idx, decision, (decode_target, decode_target_idx)


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] = []
decode_instances: list[InstanceState] = []
# Session affinity is namespace-isolated: combined-mode and pd-sep mode index
# different instance lists, so a shared dict could mis-route after a mode switch.
session_affinity_combined: dict[str, int] = {}
session_affinity_prefill: dict[str, int] = {}
# Backwards-compat alias used by /stats etc.
session_affinity = session_affinity_combined
is_pd_sep = False
_breakdown_log: list[dict] = []
_worker_state_log: list[dict] = []


async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.Event):
    for inst in instances:
        if inst.bootstrap_port is None:
            continue
        while True:
            try:
                await inst.client.get("/health")
            except Exception:
                await asyncio.sleep(1)
                continue
            parsed = urllib.parse.urlparse(str(inst.client.base_url))
            url = f"http://{parsed.hostname}:{inst.bootstrap_port}/query"
            resp = await inst.client.get(url)
            resp.raise_for_status()
            data = resp.json()
            for dp_rank, dp_entry in data.items():
                inst.engine_id[int(dp_rank)] = dp_entry["engine_id"]
            inst.dp_size = len(data)
            print(f"Inited {inst.url} engine_ids={inst.engine_id}")
            break
    ready.set()


async def _fetch_vllm_inflight(inst: "InstanceState") -> tuple[int, int] | None:
    """Read vLLM's truth: (num_running, num_waiting). Returns None on failure."""
    try:
        resp = await asyncio.wait_for(inst.client.get("/metrics"), timeout=5.0)
        if resp.status_code != 200:
            return None
        text = resp.text
    except Exception:
        return None
    running = 0
    waiting = 0
    for line in text.splitlines():
        if line.startswith("vllm:num_requests_running"):
            try:
                running = int(float(line.split()[-1]))
            except (ValueError, IndexError):
                pass
        elif line.startswith("vllm:num_requests_waiting"):
            try:
                waiting = int(float(line.split()[-1]))
            except (ValueError, IndexError):
                pass
    return running, waiting


def _engine_state_read_all(uri: str, max_age_s: float = 2.0) -> dict:
    """P2 reader (inlined; mirrors engine_state.StateReader). Returns
    {engine_id: state}, dropping records older than max_age_s."""
    now = _time.time()
    out: dict = {}
    try:
        if uri.startswith("file://"):
            import glob
            d = uri[len("file://"):]
            for p in glob.glob(os.path.join(d, "*.json")):
                try:
                    s = json.load(open(p))
                except Exception:
                    continue
                if now - s.get("ts", 0) <= max_age_s:
                    out[s.get("engine_id", os.path.basename(p)[:-5])] = s
        elif uri.startswith("redis://"):
            import redis
            r = redis.Redis.from_url(uri)
            for k in r.scan_iter("engine_state:*"):
                v = r.get(k)
                if not v:
                    continue
                s = json.loads(v)
                if now - s.get("ts", 0) <= max_age_s:
                    out[s.get("engine_id")] = s
    except Exception:
        pass
    return out


async def _engine_state_poll_loop():
    """P2: poll the engine-state feed and attach real_state to each instance.
    Instance i is keyed engine_{i} (matches AGENTIC_WORKER_ID in the launcher).
    """
    uri = SETTINGS.engine_state_uri
    if not uri:
        return
    period = max(0.01, SETTINGS.engine_state_period_ms / 1000.0)
    insts = combined_instances or (prefill_instances + decode_instances)
    print(f"[engine-state] polling {uri} every {period*1000:.0f}ms for {len(insts)} instances")
    while True:
        try:
            await asyncio.sleep(period)
        except asyncio.CancelledError:
            return
        states = await asyncio.to_thread(_engine_state_read_all, uri)
        for i, inst in enumerate(insts):
            inst.real_state = states.get(f"engine_{i}")


async def _reconcile_loop():
    """Periodic shadow-state reconciliation against vLLM /metrics truth.

    The proxy maintains shadow counters (num_requests, ongoing_tokens,
    pending_prefill_tokens, ongoing_decode_tokens) by incrementing in
    `_handle_local_request` and decrementing in the generator's finally
    block. When the generator never enters (client disconnect between
    StreamingResponse construction and Starlette starting iteration, or
    Starlette failing before iteration), the decrement never fires and
    the counter stays elevated forever. Over a long run the shadow
    accumulates "phantom" load that biases routing decisions away from
    the affected instance.

    Two-pass fix:

    1. Clamp negatives (defensive; rare in practice).
    2. Sample vLLM's actual num_running + num_waiting via /metrics. If
       the proxy's num_requests has been *higher* than vLLM's truth for
       two consecutive cycles, reconcile downward to vLLM's count.
       Two-cycle persistence avoids correcting transient mismatches
       (e.g., proxy just incremented but vLLM hasn't scheduled the
       request yet).

    Cycle period: 30 s. Two-cycle persistence threshold: 60 s of stable
    drift before correction.
    """
    prev_phantom: dict[str, int] = {}
    while True:
        try:
            await asyncio.sleep(30)
        except asyncio.CancelledError:
            return
        for inst in combined_instances + prefill_instances + decode_instances:
            # Pass 1: clamp negatives (cheap, always do).
            if inst.ongoing_tokens < 0:
                inst.ongoing_tokens = 0
            if inst.ongoing_decode_tokens < 0:
                inst.ongoing_decode_tokens = 0
            if inst.pending_prefill_tokens < 0:
                inst.pending_prefill_tokens = 0
            if inst.num_requests < 0:
                inst.num_requests = 0
            if inst.active_p_offloads < 0:
                inst.active_p_offloads = 0

            # Pass 2: detect phantom positives by polling vLLM truth.
            metrics = await _fetch_vllm_inflight(inst)
            if metrics is None:
                continue
            running, waiting = metrics
            actual_inflight = running + waiting
            phantom = inst.num_requests - actual_inflight
            prev = prev_phantom.get(inst.url, 0)
            if phantom > 0 and prev > 0:
                # Drift held across two consecutive cycles (~60 s).
                # Reconcile shadow to vLLM's truth.
                old_num = inst.num_requests
                inst.num_requests = actual_inflight
                if actual_inflight == 0:
                    # No requests in flight; zero all per-request counters.
                    inst.ongoing_tokens = 0
                    inst.ongoing_decode_tokens = 0
                    inst.pending_prefill_tokens = 0
                print(
                    f"[reconcile] {inst.url}: phantom drift "
                    f"num_requests {old_num} -> {actual_inflight} "
                    f"(vllm running={running} waiting={waiting})"
                )
            prev_phantom[inst.url] = phantom


def _verify_vllm_patch():
    """Startup self-check for patches/0001-fix-kv-transfer-abort-race.patch.

    The patch turns an `assert req_id in self.requests` into a soft warn so
    that engines do not crash on the KV-transfer abort race (see REPORT
    §3.x). If somebody upgrades vLLM without re-applying the patch, the
    assert returns and elastic mode dies under load. Print a loud warning
    so we catch the regression before the first HEAVY request.
    """
    try:
        import inspect
        from vllm.v1.core.sched.scheduler import Scheduler
        src = inspect.getsource(Scheduler)
        if "assert req_id in self.requests" in src:
            print("WARNING: vLLM scheduler still contains the unpatched "
                  "`assert req_id in self.requests` line; expect engine "
                  "death on KV-transfer abort race. Apply "
                  "patches/0001-fix-kv-transfer-abort-race.patch.")
        else:
            print("vLLM patch self-check: kv-transfer-abort assert is patched.")
    except Exception as exc:
        print(f"vLLM patch self-check skipped: {exc!r}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    global is_pd_sep
    app.state.ready = asyncio.Event()

    _verify_vllm_patch()

    reconcile_task = asyncio.create_task(_reconcile_loop())
    engine_state_task = asyncio.create_task(_engine_state_poll_loop())

    if global_args.combined:
        is_pd_sep = False
        bp_list = [int(p) for p in global_args.bootstrap_ports.split(",") if p.strip()] if global_args.bootstrap_ports else []
        for i, url in enumerate(global_args.combined):
            bp = bp_list[i] if i < len(bp_list) else None
            combined_instances.append(InstanceState(url, bp))

        # Bootstrap combined instances for offload (need engine_ids for KV transfer)
        policy = getattr(global_args, 'policy', 'linear')
        # Mooncake-based modes still need bootstrap discovery; NIXL uses
        # its own UCX side-channel and doesn't go through our proxy
        # bootstrap path. With --connector-type=nixl, v3 also skips bootstrap.
        needs_bootstrap = (
            global_args.offload
            or (policy in ("unified_v2", "unified_v3", "unified_kv_both")
                and getattr(global_args, 'connector_type', 'mooncake') == 'mooncake')
        )
        if needs_bootstrap and bp_list:
            await init_prefill_bootstrap(combined_instances, app.state.ready)
        elif needs_bootstrap and not bp_list:
            raise RuntimeError(
                f"--policy {policy} requires --bootstrap-ports for KV transfer; "
                "got empty bootstrap list."
            )
        else:
            app.state.ready.set()

        policy = getattr(global_args, 'policy', 'linear')
        print(f"Combined mode: {len(combined_instances)} instances, policy={policy}, offload={'ON' if global_args.offload else 'OFF'}")
    else:
        is_pd_sep = True
        for url, bp in global_args.prefill:
            prefill_instances.append(InstanceState(url, bp))
        for url in global_args.decode:
            decode_instances.append(InstanceState(url))
        await init_prefill_bootstrap(prefill_instances, app.state.ready)
        print(f"PD-Sep mode: {len(prefill_instances)}P + {len(decode_instances)}D")

    yield
    reconcile_task.cancel()
    try:
        await reconcile_task
    except asyncio.CancelledError:
        pass
    for inst in combined_instances + prefill_instances + decode_instances:
        await inst.client.aclose()


app = FastAPI(lifespan=lifespan)


@app.post("/v1/completions")
async def handle_completions(request: Request):
    return await _handle(request, "/v1/completions")


@app.post("/v1/chat/completions")
async def handle_chat(request: Request):
    return await _handle(request, "/v1/chat/completions")


async def _handle(request: Request, api: str):
    if not app.state.ready.is_set():
        raise HTTPException(status_code=503, detail="Service Unavailable")

    req_data = await request.json()
    incoming_rid = request.headers.get("X-Request-Id")
    request_id = incoming_rid or str(uuid.uuid4())
    prompt = req_data.get("prompt")
    token_ids = prompt if isinstance(prompt, list) else None
    input_length = len(token_ids) if token_ids else 0
    session_id = request.headers.get("X-Session-Id")

    headers = {"X-Request-Id": request_id}
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    if is_pd_sep:
        return await _handle_pd_sep(api, req_data, request_id, token_ids,
                                     input_length, session_id, headers)
    else:
        return await _handle_combined(api, req_data, token_ids,
                                       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, *, _pre_reserved: bool = False):
    breakdown.setdefault("route_class", "LOCAL")
    breakdown.setdefault("routed_to", chosen.url)
    # Skip reservation when called from _handle_combined (it already reserved
    # synchronously to close the picker→await race). When called directly
    # from non-combined paths (PD-Sep, offload), reserve here for safety.
    if not _pre_reserved:
        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()
                                breakdown["t_first_token_unix"] = _time.time()
                                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["t_done_unix"] = _time.time()
            _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):
    """Route a /v1/* request among combined (PD-colocated) instances.

    --policy options:
      linear:  cache_hit-aware load score + sticky session affinity.
      lmetric: P_tokens * BS (LMetric, OSDI'26). No session affinity.
      unified: hybrid — stick to affinity instance when cache_ratio > 0.5
               and it is not overloaded; otherwise fall back to LMetric
               with a multi-key tie-breaker.

    PD-sep offload / PUSH migration is retired (see REPORT.md §3.9 and
    commits 4c583f2 / cc6e562: relaxed-gate and forced-migration variants
    both regressed E2E tail). Re-enabling requires a new transfer mechanism.
    """
    policy = getattr(global_args, 'policy', 'linear')
    t_decision_unix = _time.time()
    request_id = headers.get("X-Request-Id", "")
    breakdown: dict = {
        "request_id": request_id,
        "session_id": session_id,
        "input_length": input_length,
        "t_proxy_recv": _time.monotonic(),
        "t_decision_unix": t_decision_unix,
        "policy": policy,
    }

    pre_decision_workers = snapshot_workers(
        combined_instances, token_ids, input_length)

    pd_sep_v2: tuple[InstanceState, int] | None = None
    if policy == "lmetric":
        chosen, best_idx = pick_instance_lmetric(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
    elif policy == "load_only":
        chosen, best_idx = pick_instance_load_only(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
    elif policy == "sticky":
        chosen, best_idx = pick_instance_sticky(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
    elif policy in ("unified", "unified_kv_both", "unified_nixl_both"):
        # unified_kv_both: same picker as `unified`, but the vLLMs are
        # launched in kv_role=kv_both with MooncakeConnector. Use this
        # as an isolation control for `unified_v2` so the v2-vs-v1 gap
        # reflects only the PD-sep branch, not the kv_both always-on
        # overhead.
        # unified_nixl_both: identical to unified_kv_both but with
        # NixlConnector at the vLLM layer. Used to attribute the
        # kv_both overhead to either Mooncake-specific code or a
        # generic v1-connector cost.
        chosen, best_idx, decision = pick_instance_unified_hybrid(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
        breakdown.update(decision)
        if session_id:
            session_affinity_combined[session_id] = best_idx
    elif policy == "unified_v2":
        chosen, best_idx, decision, pd_sep_v2 = pick_instance_unified_v2(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
        breakdown.update(decision)
        if session_id:
            session_affinity_combined[session_id] = best_idx
    elif policy == "unified_v3":
        # v3: prefill on affinity (cache reuse), decode migrated to a
        # low-load target. KV flows prefill_host → decode_target.
        # Reuses _handle_combined_pd_sep_v2 with src=prefill_host,
        # dst=decode_target (the handler is direction-agnostic).
        chosen, best_idx, decision, pd_sep_v2 = pick_instance_unified_v3(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)
        breakdown.update(decision)
        if session_id:
            if pd_sep_v2 is not None and SETTINGS.v3_rotate_affinity:
                # Migration + rotation: redirect next turn to decode_target,
                # assuming KV will live there. (Empirically wrong — see
                # cache_miss_audit. Keep behind a flag.)
                _decode_target_inst, decode_target_idx = pd_sep_v2
                session_affinity_combined[session_id] = decode_target_idx
            else:
                # No rotation: keep affinity on prefill_host (where the prefix
                # cache lives). This is the empirically correct choice.
                session_affinity_combined[session_id] = best_idx
    else:  # linear (default)
        chosen, best_idx = pick_instance(
            combined_instances, token_ids, session_id, input_length,
            session_affinity_combined)

    chosen_snap = pre_decision_workers[best_idx]
    cache_hit = chosen_snap["cache_hit"]
    estimated_new = chosen_snap["new_prefill"]
    breakdown.update({
        "cache_hit": cache_hit,
        "estimated_new_tokens": estimated_new,
        "route_class": "LOCAL" if pd_sep_v2 is None else "PD_SEP_V2",
        "routed_to": chosen.url,
        "chosen_idx": best_idx,
        "candidate_scores": pre_decision_workers,
        "chosen_score_linear": chosen_snap["score_linear"],
        "chosen_score_lmetric": chosen_snap["score_lmetric"],
    })

    _worker_state_log.append({
        "t_decision_unix": t_decision_unix,
        "request_id": request_id,
        "session_id": session_id,
        "policy": policy,
        "chosen_idx": best_idx,
        "v2_pd_sep": pd_sep_v2 is not None,
        "workers": pre_decision_workers,
    })

    if pd_sep_v2 is not None:
        # Handler contract: first arg = prefill source (does same-worker
        # prefill with do_remote_decode=True, max_tokens=1), second arg =
        # decode target (does do_remote_prefill=True, pulls KV via
        # Mooncake, decodes).
        #
        # v2 contract:  pd_sep_v2 = (src_inst, src_idx); chosen = decode
        #               → src does prefill (it has more cache), chosen decodes.
        # v3 contract:  chosen = prefill_host (affinity, has cache);
        #               pd_sep_v2 = (decode_target_inst, decode_target_idx)
        #               → chosen does prefill (cache reuse), decode_target decodes.
        if policy == "unified_v3":
            decode_target_inst, decode_target_idx = pd_sep_v2
            prefill_inst = chosen
            breakdown["v2_src_url"] = prefill_inst.url
            breakdown["v2_src_idx"] = best_idx
            breakdown["v3_decode_target_url"] = decode_target_inst.url
            breakdown["v3_decode_target_idx"] = decode_target_idx
            return await _handle_combined_pd_sep_v2(
                api, req_data, headers, token_ids, input_length,
                prefill_inst, decode_target_inst, breakdown,
                request_id=request_id)
        else:
            src_inst, src_idx = pd_sep_v2
            breakdown["v2_src_url"] = src_inst.url
            breakdown["v2_src_idx"] = src_idx
            return await _handle_combined_pd_sep_v2(
                api, req_data, headers, token_ids, input_length,
                src_inst, chosen, breakdown,
                request_id=request_id)

    # Race fix: reserve load on `chosen` BEFORE the `await` so concurrent
    # picker calls in the same asyncio event-loop tick see the updated
    # counters. Without this, two requests arriving back-to-back can both
    # pick the same "free" instance and both end up running there
    # simultaneously (observed as 8007 hotspot in A+B run).
    chosen.ongoing_tokens += input_length
    chosen.pending_prefill_tokens += estimated_new
    chosen.num_requests += 1
    breakdown.setdefault("route_class", "LOCAL")
    breakdown.setdefault("routed_to", chosen.url)
    return await _handle_local_request(
        api, req_data, headers, token_ids, input_length,
        chosen, estimated_new, breakdown, _pre_reserved=True)


async def _handle_combined_pd_sep_v2(
    api, req_data, headers, token_ids, input_length,
    src: InstanceState, dst: InstanceState, breakdown: dict,
    *, request_id: str,
):
    """Per-request PD-sep among combined instances (unified_v2 path).

    src does cached prefill (max_tokens=1) and ships KV to dst via
    Mooncake; dst pulls KV and decodes. Both instances must run in
    kv_role=kv_both with bootstrap server enabled.

    Patch 6.6: the dst streaming call uses a per-chunk read timeout
    of SETTINGS.pd_sep_xfer_timeout_s, so a stuck KV transfer fails
    the request instead of hanging for 600 s.
    """
    connector = SETTINGS.connector_type
    if connector == "mooncake" and src.bootstrap_port is None:
        raise HTTPException(
            status_code=500,
            detail=(
                "Mooncake PD-sep triggered but src instance "
                f"{src.url} has no bootstrap_port; launch with "
                "kv_role=kv_both and pass --bootstrap-ports"
            ),
        )

    # Reserve load on both endpoints.
    src.ongoing_tokens += input_length
    src.num_requests += 1
    dst.ongoing_tokens += input_length
    dst.num_requests += 1
    src_load_held = True
    dst_load_held = True

    # ---- LAYERWISE write-mode (opt-in EAR_WRITE_MODE=1, mooncake only) ------
    # Dispatch src prefill and dst decode CONCURRENTLY so the dst handshake
    # reaches src during its prefill, letting the layer-wise connector push KV
    # per-step (overlapped with prefill compute) instead of post-hoc.
    if os.environ.get("EAR_WRITE_MODE", "0") == "1" and connector == "mooncake":
        return await _handle_pdsep_v2_write_mode(
            api, req_data, headers, token_ids, input_length,
            src, dst, breakdown, request_id=request_id)

    # Build prefill kv_transfer_params per connector.
    prefill_data = req_data.copy()
    if connector == "mooncake":
        prefill_data["kv_transfer_params"] = {
            "do_remote_decode": True,
            "do_remote_prefill": False,
            "transfer_id": f"xfer-{request_id}",
        }
    else:  # nixl: src just signals it'll produce KV for remote decode
        prefill_data["kv_transfer_params"] = {"do_remote_decode": True}
    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"}

    breakdown["t_prefill_sent"] = _time.monotonic()
    breakdown["t_prefill_sent_unix"] = _time.time()
    forwarded_params: dict | None = None
    try:
        resp = await src.client.post(api, json=prefill_data, headers=p_headers)
        breakdown["t_prefill_done"] = _time.monotonic()
        breakdown["t_prefill_done_unix"] = _time.time()
        resp.raise_for_status()
        if connector == "nixl":
            # Nixl populates kv_transfer_params in the response body with
            # remote_block_ids / remote_engine_id / remote_host / remote_port.
            # We must read the body BEFORE aclose.
            src_resp_json = resp.json()
            forwarded_params = src_resp_json.get("kv_transfer_params")
            if not forwarded_params or not forwarded_params.get("remote_block_ids"):
                raise HTTPException(
                    status_code=502,
                    detail=f"Nixl src returned no remote_block_ids: {forwarded_params}",
                )
        await resp.aclose()
        src.record_prefix(token_ids)
    except Exception as e:
        breakdown["t_prefill_done"] = _time.monotonic()
        breakdown["t_prefill_done_unix"] = _time.time()
        breakdown["prefill_error"] = True
        breakdown["error_detail"] = repr(e)[:300]
        _breakdown_log.append(breakdown)
        # Release reservations on failure. Clear load_held flags so the
        # finally block below does not double-decrement (CRITICAL audit #1).
        if src_load_held:
            src.ongoing_tokens -= input_length
            src.num_requests -= 1
            src_load_held = False
        if dst_load_held:
            dst.ongoing_tokens -= input_length
            dst.num_requests -= 1
            dst_load_held = False
        raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
    finally:
        if src_load_held:
            src.ongoing_tokens -= input_length
            src.num_requests -= 1
            src_load_held = False

    decode_data = req_data.copy()
    if connector == "mooncake":
        parsed = urllib.parse.urlparse(str(src.client.base_url))
        bootstrap_addr = f"http://{parsed.hostname}:{src.bootstrap_port}"
        decode_data["kv_transfer_params"] = {
            "do_remote_decode": False,
            "do_remote_prefill": True,
            "remote_bootstrap_addr": bootstrap_addr,
            "remote_engine_id": src.engine_id.get(0, ""),
            "transfer_id": f"xfer-{request_id}",
        }
    else:  # nixl: forward what src returned
        decode_data["kv_transfer_params"] = forwarded_params

    breakdown["t_decode_sent"] = _time.monotonic()
    breakdown["t_decode_sent_unix"] = _time.time()

    xfer_timeout = httpx.Timeout(
        connect=10.0,
        read=SETTINGS.pd_sep_xfer_timeout_s,
        write=10.0,
        pool=10.0,
    )

    async def generate():
        nonlocal dst_load_held
        first_token = True
        sse_buffer = ""
        output_token_ids: list[int] = []
        try:
            async with dst.client.stream(
                "POST", api, json=decode_data, headers=headers,
                timeout=xfer_timeout,
            ) 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()
                        breakdown["t_first_token_unix"] = _time.time()
                        first_token = False
                    yield chunk
            dst.record_prefix(_realized_tokens(token_ids, output_token_ids))
        finally:
            breakdown["t_done"] = _time.monotonic()
            breakdown["t_done_unix"] = _time.time()
            if dst_load_held:
                dst.ongoing_tokens -= input_length
                dst.num_requests -= 1
                dst_load_held = False
            _breakdown_log.append(breakdown)

    return StreamingResponse(generate(), media_type="text/event-stream")


async def _handle_pdsep_v2_write_mode(
    api, req_data, headers, token_ids, input_length,
    src: InstanceState, dst: InstanceState, breakdown: dict,
    *, request_id: str,
):
    """Write-mode v3 (mooncake + MOONCAKE_LAYERWISE): dispatch src prefill and
    dst decode CONCURRENTLY so the dst handshake reaches src during prefill and
    KV is pushed per-layer (overlapped) instead of post-hoc. Caller has already
    reserved load on both src and dst; we release it in generate()'s finally.
    """
    parsed = urllib.parse.urlparse(str(src.client.base_url))
    bootstrap_addr = f"http://{parsed.hostname}:{src.bootstrap_port}"
    tid = f"xfer-{request_id}"

    prefill_data = req_data.copy()
    prefill_data["kv_transfer_params"] = {
        "do_remote_decode": True, "do_remote_prefill": False, "transfer_id": tid}
    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"}

    decode_data = req_data.copy()
    decode_data["kv_transfer_params"] = {
        "do_remote_decode": False, "do_remote_prefill": True,
        "remote_bootstrap_addr": bootstrap_addr,
        "remote_engine_id": src.engine_id.get(0, ""),
        "transfer_id": tid}

    xfer_timeout = httpx.Timeout(
        connect=10.0, read=SETTINGS.pd_sep_xfer_timeout_s, write=10.0, pool=10.0)

    breakdown["write_mode"] = True
    breakdown["t_prefill_sent"] = _time.monotonic()
    breakdown["t_prefill_sent_unix"] = _time.time()
    prefill_task = asyncio.create_task(
        src.client.post(api, json=prefill_data, headers=p_headers))
    breakdown["t_decode_sent"] = _time.monotonic()
    breakdown["t_decode_sent_unix"] = _time.time()

    async def generate():
        first_token = True
        sse_buffer = ""
        output_token_ids: list[int] = []
        try:
            async with dst.client.stream(
                "POST", api, json=decode_data, headers=headers,
                timeout=xfer_timeout,
            ) 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()
                        breakdown["t_first_token_unix"] = _time.time()
                        first_token = False
                    yield chunk
            dst.record_prefix(_realized_tokens(token_ids, output_token_ids))
        finally:
            breakdown["t_done"] = _time.monotonic()
            breakdown["t_done_unix"] = _time.time()
            try:
                presp = await prefill_task
                breakdown["t_prefill_done"] = _time.monotonic()
                breakdown["t_prefill_done_unix"] = _time.time()
                presp.raise_for_status()
                await presp.aclose()
                src.record_prefix(token_ids)
            except Exception as e:
                breakdown["prefill_error"] = True
                breakdown["error_detail"] = repr(e)[:300]
            src.ongoing_tokens -= input_length
            src.num_requests -= 1
            dst.ongoing_tokens -= input_length
            dst.num_requests -= 1
            _breakdown_log.append(breakdown)

    return StreamingResponse(generate(), media_type="text/event-stream")


async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
                          session_id, headers):
    """PD-Sep mode with per-stage breakdown profiling."""
    t_decision_unix = _time.time()
    breakdown = {
        "request_id": request_id,
        "session_id": session_id,
        "input_length": input_length,
        "t_proxy_recv": _time.monotonic(),
        "t_decision_unix": t_decision_unix,
        "policy": "pd_sep",
    }

    pre_decision_p = snapshot_workers(prefill_instances, token_ids, input_length)
    pre_decision_d = snapshot_workers(decode_instances, token_ids, input_length)

    p_inst, p_idx = pick_instance(prefill_instances, token_ids, session_id,
                                   input_length, session_affinity_prefill)
    d_idx = min(range(len(decode_instances)),
                 key=lambda i: decode_instances[i].ongoing_tokens)
    d_inst = decode_instances[d_idx]
    breakdown["p_inst"] = p_inst.url
    breakdown["d_inst"] = d_inst.url
    breakdown["candidate_scores_prefill"] = pre_decision_p
    breakdown["candidate_scores_decode"] = pre_decision_d
    breakdown["chosen_p_idx"] = p_idx
    breakdown["chosen_d_idx"] = d_idx

    _worker_state_log.append({
        "t_decision_unix": t_decision_unix,
        "request_id": request_id,
        "session_id": session_id,
        "policy": "pd_sep",
        "chosen_p_idx": p_idx,
        "chosen_d_idx": d_idx,
        "workers_prefill": pre_decision_p,
        "workers_decode": pre_decision_d,
    })

    prefill_data = req_data.copy()
    prefill_data["kv_transfer_params"] = {
        "do_remote_decode": True, "do_remote_prefill": False,
        "transfer_id": f"xfer-{request_id}",
    }
    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"}

    p_inst.ongoing_tokens += input_length
    p_inst.num_requests += 1
    breakdown["t_prefill_sent"] = _time.monotonic()
    breakdown["t_prefill_sent_unix"] = _time.time()

    try:
        resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
        breakdown["t_prefill_done"] = _time.monotonic()
        breakdown["t_prefill_done_unix"] = _time.time()
        resp.raise_for_status()
        await resp.aclose()
        p_inst.record_prefix(token_ids)
    except Exception as e:
        breakdown["t_prefill_done"] = _time.monotonic()
        breakdown["t_prefill_done_unix"] = _time.time()
        breakdown["prefill_error"] = True
        _breakdown_log.append(breakdown)
        raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
    finally:
        p_inst.ongoing_tokens -= input_length
        p_inst.num_requests -= 1

    # Send decode
    d_inst.ongoing_tokens += input_length
    d_inst.num_requests += 1
    parsed = urllib.parse.urlparse(str(p_inst.client.base_url))
    bootstrap_addr = f"http://{parsed.hostname}:{p_inst.bootstrap_port}"

    decode_data = req_data.copy()
    decode_data["kv_transfer_params"] = {
        "do_remote_decode": False, "do_remote_prefill": True,
        "remote_bootstrap_addr": bootstrap_addr,
        "remote_engine_id": p_inst.engine_id.get(0, ""),
        "transfer_id": f"xfer-{request_id}",
    }

    breakdown["t_decode_sent"] = _time.monotonic()
    breakdown["t_decode_sent_unix"] = _time.time()

    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()
                        breakdown["t_first_token_unix"] = _time.time()
                        first_token = False
                    yield chunk
            d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
        finally:
            breakdown["t_done"] = _time.monotonic()
            breakdown["t_done_unix"] = _time.time()
            d_inst.ongoing_tokens -= input_length
            d_inst.num_requests -= 1
            _breakdown_log.append(breakdown)

    return StreamingResponse(generate(), media_type="text/event-stream")


@app.get("/breakdown")
async def get_breakdown():
    """Return per-request breakdown data for analysis."""
    return _breakdown_log


@app.get("/worker_state")
async def get_worker_state():
    """Return per-decision worker-state snapshot log (one entry per route decision)."""
    return _worker_state_log


@app.get("/worker_state/latest")
async def get_worker_state_latest():
    """Return current per-worker state snapshot without recording it."""
    if combined_instances:
        return {
            "t_unix": _time.time(),
            "mode": "combined",
            "workers": snapshot_workers(combined_instances),
        }
    return {
        "t_unix": _time.time(),
        "mode": "pd_sep",
        "workers_prefill": snapshot_workers(prefill_instances),
        "workers_decode": snapshot_workers(decode_instances),
    }


@app.get("/stats")
async def get_stats():
    """Return per-instance live state for debugging."""
    instances = combined_instances or prefill_instances + decode_instances
    return [{
        "url": inst.url,
        "role": "combined",
        "ongoing_tokens": inst.ongoing_tokens,
        "pending_prefill_tokens": inst.pending_prefill_tokens,
        "ongoing_decode_tokens": inst.ongoing_decode_tokens,
        "num_requests": inst.num_requests,
        "active_p_offloads": inst.active_p_offloads,
        "cached_blocks": len(inst.cached_blocks),
    } for inst in instances]


def parse_args():
    p = argparse.ArgumentParser(description="Unified cache-aware global scheduler")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument("--host", type=str, default="0.0.0.0")
    p.add_argument("--combined", nargs="+", help="Combined mode: list of instance URLs")
    p.add_argument("--prefill", nargs="+", action="append", dest="prefill_raw",
                   help="PD-Sep prefill: URL [bootstrap_port]")
    p.add_argument("--decode", nargs=1, action="append", dest="decode_raw",
                   help="PD-Sep decode: URL")
    p.add_argument("--heavy-threshold", type=int, default=20000,
                   help="New tokens threshold for HEAVY classification (adaptive offload)")
    p.add_argument("--offload", action="store_true",
                   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", "load_only", "sticky",
                             "unified", "unified_kv_both",
                             "unified_nixl_both", "unified_v2",
                             "unified_v3"],
                   help="Routing policy: linear (cache-aware), lmetric (P_tokens × BS), "
                        "load_only (B3 control: pure min-num_requests), "
                        "sticky (B3 control: hard session affinity), "
                        "unified (hybrid affinity + LMetric fallback), "
                        "unified_kv_both (unified picker on kv_both Mooncake "
                        "vLLMs; isolation control for unified_v2), "
                        "unified_nixl_both (same as unified_kv_both but using "
                        "NixlConnector instead of MooncakeConnector; isolates "
                        "connector implementation from policy effect), "
                        "or unified_v2 (unified + selective per-request PD-sep "
                        "via Mooncake; requires --bootstrap-ports and "
                        "kv_role=kv_both vLLM launch)")
    p.add_argument("--v3-rotate-affinity", type=int, default=1, choices=[0,1],
                   help="unified_v3 only: 1 = rotate session affinity to decode_target "
                        "after migration (original behavior, empirically loses prefix cache); "
                        "0 = keep affinity on prefill_host so next turn hits its cache.")
    p.add_argument("--connector-type", type=str, default="mooncake",
                   choices=["mooncake", "nixl"],
                   help="PD-sep handshake protocol. 'mooncake' uses pre-baked engine_id"
                        " + bootstrap_addr (requires --bootstrap-ports). 'nixl' uses"
                        " response-forward (src returns kv_transfer_params, proxy"
                        " relays to dst; Nixl/UCX auto-picks NVLink intra-node).")
    p.add_argument("--v3-prefer-cache-target", type=int, default=1, choices=[0,1],
                   help="Mechanism B: unified_v3 picks decode_target with the most"
                        " prefix cache among low-load candidates (default 1). Set 0"
                        " to fall back to pure-load tie-break (cache-blind).")
    p.add_argument("--lmetric-decode-weight", type=float, default=0.0,
                   help="Direction B: LMetric fallback adds this × ongoing_decode_tokens"
                        " to the queue-depth score, so hosts with heavy decode load get"
                        " penalised. 0 = original behavior; 0.01 is a reasonable start.")
    p.add_argument("--overload-factor", type=float, default=2.0,
                   help="Break session affinity when instance load > factor * avg")
    p.add_argument("--engine-state-uri", type=str, default="",
                   help="P2: real engine-state feed for migration target "
                        "selection (file:///dev/shm/... or redis://...). "
                        "Empty=disabled (shadow counters only).")
    # The four flags below are accepted for bench.sh backward compatibility but
    # have no effect after the PD-sep offload path was retired (REPORT §3.9,
    # commits 4c583f2 / cc6e562). Removing them would break scripts/bench.sh and
    # scripts/legacy/*.sh which still pass them through.
    p.add_argument("--max-offload-inflight", type=int, default=4,
                   help="[DEPRECATED] PUSH offload retired; no effect")
    p.add_argument("--offload-mode", type=str, default="cached_prefill",
                   choices=["direct_read", "cached_prefill"],
                   help="[DEPRECATED] PUSH offload retired; no effect")
    p.add_argument("--cache-gate-ratio", type=float, default=0.0,
                   help="[DEPRECATED] PUSH offload retired; no effect")
    p.add_argument("--decode-iteration-s", type=float, default=0.05,
                   help="[DEPRECATED] PUSH offload retired; no effect")
    args = p.parse_args()

    args.prefill = []
    if args.prefill_raw:
        for entry in args.prefill_raw:
            url = entry[0]
            bp = int(entry[1]) if len(entry) > 1 and entry[1].lower() != "none" else None
            args.prefill.append((url, bp))
    args.decode = [e[0] for e in (args.decode_raw or [])]

    if not args.combined and not args.prefill:
        p.error("Must specify either --combined or --prefill/--decode")
    return args


if __name__ == "__main__":
    global_args = parse_args()
    SETTINGS.heavy_threshold = global_args.heavy_threshold
    SETTINGS.overload_factor = global_args.overload_factor
    SETTINGS.engine_state_uri = getattr(global_args, 'engine_state_uri', '') or ''
    SETTINGS.max_offload_inflight = global_args.max_offload_inflight
    SETTINGS.cache_gate_ratio = global_args.cache_gate_ratio
    SETTINGS.decode_iteration_s = getattr(global_args, 'decode_iteration_s', 0.05)
    SETTINGS.v3_rotate_affinity = bool(getattr(global_args, 'v3_rotate_affinity', 1))
    SETTINGS.connector_type = getattr(global_args, 'connector_type', 'mooncake')
    SETTINGS.v3_prefer_cache_target = bool(getattr(global_args, 'v3_prefer_cache_target', 1))
    SETTINGS.lmetric_decode_weight = float(getattr(global_args, 'lmetric_decode_weight', 0.0))
    print("SETTINGS: throughput=%.0f rdma_overhead=%.2f offload=%s v3_rotate_affinity=%s "
          "connector_type=%s v3_prefer_cache_target=%s lmetric_decode_weight=%.3f" % (
        SETTINGS.prefill_throughput, SETTINGS.rdma_overhead_s,
        getattr(global_args, 'offload', False),
        SETTINGS.v3_rotate_affinity,
        SETTINGS.connector_type,
        SETTINGS.v3_prefer_cache_target,
        SETTINGS.lmetric_decode_weight))
    uvicorn.run(app, host=global_args.host, port=global_args.port)
