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

PS experiments + H4 cache-gate + GPU profiling + Mooncake elif→if fix

Browse Source 
Experiments run:
- Phase 0: kv_both has zero idle overhead (TPOT +1.3%, noise)
- PS V1 (cold prefill): REJECTED — PS always slower than cached C
- PS V1+flexD: 92.5% OK, HEAVY TTFT 7.8s (baseline 5.0s) — PS bottleneck
- V2 (C_s prefill + flexible D): E2E -9% but 6 errors, RDMA bimodal
- H4 (cache-gate): 198/200 OK, GPU imbalance 4.0x→2.0x, but HEAVY_OFFLOAD
  TTFT=11.5s due to RDMA. HEAVY_COLO improved 10.5% from better balance.
- H5: Mooncake RDMA transfer R²=0.095, bimodal (0.6s or 18-30s)

Key findings:
- Mooncake lacks layerwise KV transfer → RDMA is pure sequential overhead
- 92% of HEAVY are turn-1 cold → offloading cold requests always loses
- GPU balance improvement from routing IS real (-10.5% HEAVY_COLO TTFT)
- RDMA transfer negates the routing benefit for offloaded requests

Code changes:
- bench.sh: add GPU timeline monitoring (gpu_monitor.sh during benchmark)
- cache_aware_proxy.py: H4 cache-gate, flexible D, PS routing
- mooncake_connector.py: elif→if fix (allow dual prefill+decode flags)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-23 02:14:37 +08:00
parent 098d86385a
commit 3bc37cc6d5
11 changed files with 1095 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

121
scripts/cache_aware_proxy.py
View File

@@ -43,6 +43,7 @@ class InstanceState:
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
self.cached_blocks: set[int] = set()
@@ -72,14 +73,28 @@ class InstanceState:
_inst_cumulative_tokens: list[int] = []
def _p_offload_penalty(inst: InstanceState) -> int:
"""Penalty for instances currently doing P-role offloaded prefills.
When an instance is busy with offloaded HEAVY prefills for other
instances, we want to steer WARM/MEDIUM requests away from it so
its GPU is dedicated to prefill (soft PD separation).
"""
if inst.active_p_offloads <= 0:
return 0
return inst.active_p_offloads * HEAVY_THRESHOLD
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
(ongoing_tokens > 2x average), in which case pick least-loaded.
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.
"""
global _inst_cumulative_tokens
if not _inst_cumulative_tokens:
@@ -87,22 +102,19 @@ def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
avg_load = max(sum(i.ongoing_tokens for i in instances) / len(instances), 1.0)
# Session affinity for turn 2+ (with load override)
if session_id and session_id in affinity:
idx = affinity[session_id]
if idx < len(instances):
inst = instances[idx]
# Stick if not overloaded
if inst.ongoing_tokens <= avg_load * OVERLOAD_FACTOR:
if (inst.ongoing_tokens <= avg_load * OVERLOAD_FACTOR
and inst.active_p_offloads == 0):
return inst, idx
# Overloaded: fall through to score-based selection
# Score = ongoing_tokens - ALPHA * cache_hit_tokens
# Balances load (lower is better) with cache affinity (higher hit is better)
best_idx, best_score = 0, float("inf")
for i, inst in enumerate(instances):
cache_hit = inst.estimate_cache_hit(token_ids)
score = inst.ongoing_tokens - CACHE_HIT_ALPHA * cache_hit
score = (inst.ongoing_tokens + _p_offload_penalty(inst)
- CACHE_HIT_ALPHA * cache_hit)
if score < best_score:
best_score = score
best_idx = i
@@ -118,8 +130,7 @@ def pick_instance_lmetric(instances: list[InstanceState], token_ids: list[int] |
affinity: dict[str, int]) -> tuple[InstanceState, int]:
"""LMetric routing: score = P_tokens × BS (OSDI'26).
P_tokens = pending_prefill_tokens on instance + new request's uncached tokens.
BS = num_requests on instance + 1 (counting the new request).
Instances doing P-role offloads get a large penalty.
"""
avg_load = max(sum(i.ongoing_tokens for i in instances) / len(instances), 1.0)
@@ -127,14 +138,15 @@ def pick_instance_lmetric(instances: list[InstanceState], token_ids: list[int] |
idx = affinity[session_id]
if idx < len(instances):
inst = instances[idx]
if inst.ongoing_tokens <= avg_load * OVERLOAD_FACTOR:
if (inst.ongoing_tokens <= avg_load * 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)
new_prefill = max(0, input_length - cache_hit)
p_tokens = inst.pending_prefill_tokens + new_prefill
p_tokens = inst.pending_prefill_tokens + new_prefill + _p_offload_penalty(inst)
bs = inst.num_requests + 1
score = p_tokens * bs
if score < best_score:
@@ -153,9 +165,6 @@ decode_instances: list[InstanceState] = []
session_affinity: dict[str, int] = {}
is_pd_sep = False
_breakdown_log: list[dict] = []
_offload_inflight = 0 # number of currently in-flight offloaded HEAVY requests
MAX_OFFLOAD_INFLIGHT = 4 # cap concurrent offloads to prevent P overload
_p_round_robin_idx = 0 # round-robin counter for P-instance selection
async def init_prefill_bootstrap(instances: list[InstanceState], ready: asyncio.Event):
@@ -192,10 +201,13 @@ async def lifespan(app: FastAPI):
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)
if global_args.offload and bp_list:
await init_prefill_bootstrap(combined_instances, app.state.ready)
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:
@@ -250,12 +262,12 @@ async def _handle(request: Request, api: str):
async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers):
"""Combined mode with adaptive prefill offload (v2).
"""Combined mode with V2 P2P offload.
WARM/MEDIUM: route to best instance, co-located P+D (no KV transfer).
HEAVY (kv_both mode): P on least-loaded instance, KV via Mooncake, D on
session-sticky instance. Only works if instances have kv_role=kv_both.
Falls back to co-located if --no-offload or instances lack Mooncake.
HEAVY: C_s (session-sticky, has cache) does FAST prefill,
D (least-loaded C, D != C_s) pulls KV via Mooncake and decodes.
Offload only when D is meaningfully less loaded than C_s.
"""
policy = getattr(global_args, 'policy', 'linear') if global_args else 'linear'
picker = pick_instance_lmetric if policy == 'lmetric' else pick_instance
@@ -272,50 +284,45 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
"t_proxy_recv": _time.monotonic(),
}
offload_enabled = getattr(global_args, 'offload', False) if global_args else False
has_bootstrap = any(inst.bootstrap_port for inst in combined_instances)
# Elastic offload decision: offload only when it helps
# H4 cache-aware offload gate: only offload when C_s has significant cache
# Cold turn-1 HEAVY: stay co-located (no RDMA overhead)
# Cached turn-2+ HEAVY: offload to flexible D (C_s fast prefill + D decode)
offload_enabled = getattr(global_args, 'offload', False) and len(combined_instances) >= 2
use_offload = False
offload_reason = "disabled"
if estimated_new >= HEAVY_THRESHOLD and offload_enabled and has_bootstrap and len(combined_instances) >= 2:
d_inst = best_inst
p_candidates = [(i, inst) for i, inst in enumerate(combined_instances) if inst is not d_inst]
avg_load = max(sum(i.ongoing_tokens for i in combined_instances) / len(combined_instances), 1.0)
offload_reason = "offload_disabled"
# Round-robin P selection with overload skip (spreads P-role evenly)
global _offload_inflight, _p_round_robin_idx
p_inst = None
for _ in range(len(p_candidates)):
_p_round_robin_idx = (_p_round_robin_idx + 1) % len(p_candidates)
candidate = p_candidates[_p_round_robin_idx][1]
if candidate.ongoing_tokens < avg_load * OVERLOAD_FACTOR:
p_inst = candidate
break
if p_inst is None:
p_inst = min(p_candidates, key=lambda x: x[1].ongoing_tokens)[1]
if estimated_new >= HEAVY_THRESHOLD and offload_enabled:
cache_ratio = cache_hit / max(input_length, 1)
d_candidate = min((c for c in combined_instances if c is not best_inst),
key=lambda c: c.ongoing_tokens)
breakdown["cache_ratio"] = cache_ratio
if _offload_inflight >= MAX_OFFLOAD_INFLIGHT:
offload_reason = "max_concurrent_reached"
elif p_inst.ongoing_tokens >= HEAVY_THRESHOLD * 2:
offload_reason = "p_saturated"
else:
if cache_ratio >= 0.3: # at least 30% cache hit to justify RDMA offload
use_offload = True
offload_reason = "offload_accepted"
_offload_inflight += 1
offload_reason = "cached_offload_%.0f%%" % (cache_ratio * 100)
else:
offload_reason = "cold_colocated_%.0f%%" % (cache_ratio * 100)
if use_offload:
d_idx = best_idx
p_inst.ongoing_tokens += input_length # reserve immediately
# C_s does fast cached prefill, D does decode
p_inst = best_inst # session-sticky, has prefix cache
d_inst = d_candidate
d_idx = combined_instances.index(d_inst)
# Accounting: C_s only prefills estimated_new tokens (cached prefix is free)
p_inst.ongoing_tokens += input_length
p_inst.pending_prefill_tokens += estimated_new
p_inst.num_requests += 1
p_inst.active_p_offloads += 1
breakdown["route_class"] = "HEAVY_P2P"
breakdown["route_class"] = "HEAVY_OFFLOAD"
breakdown["offload_reason"] = offload_reason
breakdown["p_inst"] = p_inst.url
breakdown["d_inst"] = d_inst.url
breakdown["p_load"] = p_inst.ongoing_tokens
breakdown["d_load"] = d_inst.ongoing_tokens
# Update session affinity to D (D will have KV after this request)
if session_id:
session_affinity[session_id] = d_idx
@@ -325,8 +332,10 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
if estimated_new >= HEAVY_THRESHOLD:
breakdown["route_class"] = "HEAVY_COLO"
breakdown["offload_reason"] = offload_reason
elif estimated_new < 5000:
breakdown["route_class"] = "WARM"
else:
breakdown["route_class"] = "WARM" if estimated_new < 5000 else "MEDIUM"
breakdown["route_class"] = "MEDIUM"
inst = best_inst
breakdown["routed_to"] = inst.url
@@ -366,13 +375,15 @@ PREFILL_TIMEOUT_S = 120 # max seconds to wait for P-instance prefill
async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length,
p_inst, d_inst, breakdown):
"""HEAVY request: prefill on p_inst, KV via Mooncake, decode on d_inst.
"""HEAVY request: prefill on p_inst (C_s), KV via Mooncake, decode on d_inst (D).
On prefill timeout/failure, falls back to co-located decode on d_inst.
"""
global _offload_inflight
request_id = headers.get("X-Request-Id", "")
estimated_new = breakdown.get("estimated_new_tokens", 0)
# V2: p_inst is C_s with cache, so pending_prefill_tokens was incremented
# by estimated_new (only new tokens), not full input_length.
p_prefill_release = estimated_new
# Step 1: Await prefill on p_inst (ongoing_tokens already reserved by caller)
breakdown["t_prefill_sent"] = _time.monotonic()
@@ -405,9 +416,9 @@ async def _handle_heavy_offload(api, req_data, headers, token_ids, input_length,
finally:
# Always release P-instance resources exactly once
p_inst.ongoing_tokens -= input_length
p_inst.pending_prefill_tokens -= estimated_new
p_inst.pending_prefill_tokens -= p_prefill_release
p_inst.num_requests -= 1
_offload_inflight = max(0, _offload_inflight - 1)
p_inst.active_p_offloads = max(0, p_inst.active_p_offloads - 1)
if not prefill_ok:
# Fallback: co-located prefill+decode on d_inst (no KV transfer)
@@ -587,10 +598,12 @@ async def get_stats():
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]