proxy: real LRU for cached_blocks + shadow-state reconcile loop (M1, M5)

M1: cached_blocks was a plain set with a "trim half via list slicing"
eviction. CPython does not guarantee set iteration order, so the trim
discarded an arbitrary half of the entries — completely unlike vLLM's
LRU and a known contributor to the router's cache_hit estimate
diverging from real APC. Replace with an OrderedDict-backed LRU:
move_to_end on hits, popitem(last=False) on overflow. Capacity exposed
as CACHE_CAPACITY_BLOCKS module constant (200000 by default).

M5: streamed responses decrement load counters in their generator's
finally block. If a client disconnects before consuming the body the
generator is never entered and the decrement is lost, causing
ongoing_tokens / num_requests / pending_prefill_tokens to drift
negative under load. Add a 60s background reconcile_loop that clamps
those counters at zero as a safety net. Started in lifespan, cancelled
on shutdown. Does not replace proper vLLM exact-state syncing.

scripts/cache_aware_proxy.py

@@ -18,6 +18,7 @@ import os
 import time as _time
 import urllib.parse
 import uuid
+from collections import OrderedDict
 from contextlib import asynccontextmanager
 
 import httpx
@@ -32,6 +33,7 @@ OVERLOAD_FACTOR = 2.0    # default; overridden by --overload-factor
 MAX_OFFLOAD_INFLIGHT = 4  # cap concurrent P-role offloads
 PREFILL_THROUGHPUT = 7000  # tokens/s per GPU (from H20 measurements)
 RDMA_OVERHEAD_S = 2.0     # seconds of RDMA transfer + decode start overhead
+CACHE_CAPACITY_BLOCKS = 200000  # per-instance LRU cap on shadow cached_blocks
 
 
 class InstanceState:
@@ -49,7 +51,8 @@ class InstanceState:
         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
-        self.cached_blocks: set[int] = set()
+        # OrderedDict acts as an LRU keyed by block hash; value is unused.
+        self.cached_blocks: OrderedDict[int, None] = OrderedDict()
 
     def estimate_cache_hit(self, token_ids: list[int] | None) -> int:
         if not token_ids or len(token_ids) < BLOCK_SIZE:
@@ -58,6 +61,7 @@ class InstanceState:
         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
@@ -67,9 +71,13 @@ class InstanceState:
         if not token_ids:
             return
         for i in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):
-            self.cached_blocks.add(hash(tuple(token_ids[i:i + BLOCK_SIZE])))
-        if len(self.cached_blocks) > 200000:
-            self.cached_blocks = set(list(self.cached_blocks)[-100000:])
+            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) > CACHE_CAPACITY_BLOCKS:
+                    self.cached_blocks.popitem(last=False)
 
 
 def _p_offload_penalty(inst: InstanceState) -> int:
@@ -174,11 +182,40 @@ async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.
     ready.set()
 
 
+async def _reconcile_loop():
+    """Periodic safety net for shadow state.
+
+    StreamingResponse generators decrement load counters in their finally
+    block, but if a client disconnects before the body is consumed the
+    generator is never entered and the decrement is lost. Clamp negative
+    drift every minute so router scores stay sane. This does not replace
+    proper exact-state syncing with vLLM (see TODO.md item 6).
+    """
+    while True:
+        try:
+            await asyncio.sleep(60)
+        except asyncio.CancelledError:
+            return
+        for inst in combined_instances + prefill_instances + decode_instances:
+            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
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global is_pd_sep
     app.state.ready = asyncio.Event()
 
+    reconcile_task = asyncio.create_task(_reconcile_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 []
@@ -204,6 +241,11 @@ async def lifespan(app: FastAPI):
         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()