gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Unified-routing A+B ablation: decode-aware LMetric + v3 anti-hotspot

Browse Source 
cache_aware_proxy: add lmetric_decode_weight (decode-load penalty in the
LMetric fallback score) and a v3 anti-hotspot recent-migration penalty
(effective_load = num_req + recent-migration count over a sliding window),
preventing back-to-back migration clustering. UNIFIED_ABLATION.md documents
the A (overload_factor=1.3) + B' (decode-weight, max(num_req,1)) + RaceFix
sweep: A+B'+RaceFix reaches TTFT p90 7770ms, beating v3 PD-sep migration by
~20%. Runners/analyzer for the b3 trace replay included.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 11:52:44 +08:00
parent a2f2645fda
commit 67fcec7933
7 changed files with 984 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
scripts/cache_aware_proxy.py
View File

@@ -19,7 +19,7 @@ import os
import time as _time
import urllib.parse
import uuid
from collections import OrderedDict
from collections import OrderedDict, deque
from contextlib import asynccontextmanager
from dataclasses import dataclass
@@ -103,6 +103,20 @@ class Settings:
# 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
# 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).
@@ -187,6 +201,11 @@ class InstanceState:
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)
def estimate_cache_hit(self, token_ids: list[int] | None) -> int:
if not token_ids or len(token_ids) < BLOCK_SIZE:
@@ -417,13 +436,24 @@ def pick_instance_unified_hybrid(
decision["chosen_idx"] = a_idx
return a_inst, a_idx, decision
keys: list[tuple[int, int, int, int]] = []
# 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)
p_tokens = inst.pending_prefill_tokens + new_prefill
decode_pen = SETTINGS.lmetric_decode_weight * inst.ongoing_decode_tokens
bs = inst.num_requests
score = p_tokens * bs
score = (p_tokens + decode_pen) * max(bs, 1)
keys.append((score, new_prefill, bs, i))
best_triple = min(k[:3] for k in keys)
@@ -637,48 +667,80 @@ def pick_instance_unified_v3(
)
return prefill_host, prefill_idx, decision, None
# Gate 3: pick the lowest-load target that is materially less loaded
# than the prefill_host. Cache content irrelevant — KV ships over.
# 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 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 inst.num_requests + recent * SETTINGS.v3_recent_mig_weight
threshold_loaded = max(1,
int(prefill_host.num_requests * SETTINGS.v3_target_load_ratio))
candidates = [
(i, inst) for i, inst in enumerate(instances)
if i != prefill_idx
and inst.num_requests < threshold_loaded
and inst.num_requests <= prefill_host.num_requests - SETTINGS.v3_min_load_gap
and effective_load(inst) < threshold_loaded
and effective_load(inst) <= prefill_host.num_requests - 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"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 load. vLLM
# auto-skips transferring overlapping prefix when dst's local cache
# matches — verified in smoke_partial_transfer: 77% faster on a 33k
# prompt when dst has the prefix already.
# 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:
decode_target_idx, decode_target = min(
candidates,
key=lambda x: (-x[1].estimate_cache_hit(token_ids),
x[1].num_requests, x[1].ongoing_tokens))
effective_load(x[1]),
x[1].ongoing_tokens))
else:
decode_target_idx, decode_target = min(
candidates, key=lambda x: (x[1].num_requests, x[1].ongoing_tokens))
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"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)
@@ -987,12 +1049,16 @@ async def _handle(request: Request, api: str):
async def _handle_local_request(api, req_data, headers, token_ids, input_length,
chosen: InstanceState, estimated_new: int,
breakdown: dict):
breakdown: dict, *, _pre_reserved: bool = False):
breakdown.setdefault("route_class", "LOCAL")
breakdown.setdefault("routed_to", chosen.url)
chosen.ongoing_tokens += input_length
chosen.pending_prefill_tokens += estimated_new
chosen.num_requests += 1
# 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
@@ -1180,9 +1246,19 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
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)
chosen, estimated_new, breakdown, _pre_reserved=True)
async def _handle_combined_pd_sep_v2(
@@ -1545,6 +1621,10 @@ def parse_args():
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")
# The four flags below are accepted for bench.sh backward compatibility but
@@ -1585,11 +1665,13 @@ if __name__ == "__main__":
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" % (
"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.v3_prefer_cache_target,
SETTINGS.lmetric_decode_weight))
uvicorn.run(app, host=global_args.host, port=global_args.port)