From 020be9f44424ebc5eb02be5b5d9b01251f66720d Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sat, 23 May 2026 21:00:35 +0800 Subject: [PATCH] proxy: real LRU for cached_blocks + shadow-state reconcile loop (M1, M5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 50 +++++++++++++++++++++++++++++++++--- 1 file changed, 46 insertions(+), 4 deletions(-) diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 03f4de2..7708bb3 100644 --- a/scripts/cache_aware_proxy.py +++ b/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()